Support scalar parameters in scale_by_sm3

[TheSaiEaranti]

Problem

optax.sm3 crashes with IndexError: list index out of range on the first update() whenever the parameter pytree contains a scalar (0-d) leaf — a temperature, a scalar gain — while every other transformation in the library (adam, adagrad, rmsprop, ...) accepts the identical pytree:

params = {'w': jnp.ones((2, 3)), 'temperature': jnp.asarray(2.0)}
opt = optax.sm3(0.01)
state = opt.init(params)          # succeeds
opt.update(grads, state, params)  # IndexError: list index out of range

init builds the per-axis accumulator list with range(p.ndim) — the empty list for a 0-d leaf — and the first update then indexes v[0] inside _new_accum. Neither the docstring nor an error message surfaces any shape restriction; init succeeding and update throwing an unrelated IndexError is the worst version of the failure.

Fix

Treat a 0-d leaf as having one virtual axis, the algorithm's natural degenerate case: one accumulator covering the whole (scalar) parameter. Three range(max(ndim, 1)) sites plus making _expanded_shape rank-0-safe via slicing. Behavior for all existing shapes is unchanged (shape[axis:axis+1] equals [shape[axis]] for every rank ≥ 1).

The regression test asserts the crash is gone and the semantics are right: a scalar leaf's update equals the same value's update as a one-element vector.

Tests

test_scale_by_sm3_scalar_params fails on main with the IndexError and passes with the fix. Full transform_test.py + alias_test.py: 560 passed, 643 subtests. ruff clean.