Support scalar parameters in scale_by_sm3

optax/_src/transform.py

@@ -1198,7 +1198,9 @@ def scale_by_sm3(
   """
 
   def zeros_for_dim(p):
-    return [_zeros_like_axis(p, i) for i in range(p.ndim)]
+    # A scalar (0-d) leaf is treated as having one virtual axis so that the
+    # accumulator list is never empty.
+    return [_zeros_like_axis(p, i) for i in range(max(p.ndim, 1))]
 
   def init_fn(params):
     _reject_complex(params)
@@ -1210,7 +1212,7 @@ def _expanded_shape(shape, axis):
     # Replaces a `shape` of [M, N, K] with 1 in all dimensions except for i.
     # For eg: i = 1 returns [1, N, 1].
     rank = len(shape)
-    return [1] * axis + [shape[axis]] + [1] * (rank - axis - 1)
+    return [1] * axis + list(shape[axis : axis + 1]) + [1] * (rank - axis - 1)
 
   def _new_accum(g, v):
     coeffs = ((1.0 - b2) if b2 != 1.0 else 1.0, b2)
@@ -1233,7 +1235,8 @@ def update_fn(updates, state, params=None):
 
     def f(g, v):
       return [
-          jnp.reshape(v[i], _expanded_shape(g.shape, i)) for i in range(g.ndim)
+          jnp.reshape(v[i], _expanded_shape(g.shape, i))
+          for i in range(max(g.ndim, 1))
       ]
 
     mu = jax.tree.map(f, updates, state.mu)
@@ -1243,7 +1246,9 @@ def f(g, v):
     )
     up = jax.tree.map(lambda g, a: g * a, updates, accum_inv_sqrt)
     nu = optax.tree.update_moment(up, state.nu, b1, 1)
-    mu = jax.tree.map(lambda g: [_new_mu(g, i) for i in range(g.ndim)], accum)
+    mu = jax.tree.map(
+        lambda g: [_new_mu(g, i) for i in range(max(g.ndim, 1))], accum
+    )
 
     return nu, ScaleBySM3State(mu=mu, nu=nu)
 

optax/_src/transform_test.py

@@ -73,6 +73,31 @@ def test_scalers(self, scaler_constr):
     test_utils.assert_tree_all_finite((params, updates, state))
     test_utils.assert_trees_all_equal_shapes(params, updates)
 
+  def test_scale_by_sm3_scalar_params(self):
+    """sm3 must accept pytrees containing scalar (0-d) leaves.
+
+    The per-axis accumulator list used to be empty for a 0-d leaf, so the
+    first update raised IndexError, while every other transformation accepts
+    the same pytree.
+    """
+    params = {'matrix': jnp.ones((2, 3)), 'temperature': jnp.asarray(2.0)}
+    updates = jax.tree.map(lambda p: 0.5 * jnp.ones_like(p), params)
+
+    tx = transform.scale_by_sm3()
+    state = tx.init(params)
+    scaled, _ = tx.update(updates, state, params)
+    self.assertEqual(scaled['temperature'].shape, ())
+
+    # A scalar leaf must behave exactly like the same value as a 1-element
+    # vector, the natural degenerate case of the algorithm.
+    vec_params = {'temperature': jnp.asarray([2.0])}
+    vec_updates = {'temperature': jnp.asarray([0.5])}
+    vec_state = tx.init(vec_params)
+    vec_scaled, _ = tx.update(vec_updates, vec_state, vec_params)
+    test_utils.assert_trees_all_close(
+        scaled['temperature'], vec_scaled['temperature'][0], rtol=1e-6
+    )
+
   def test_apply_every(self):
     # The frequency of the application of sgd
     k = 4