Support negative axes in microbatching

optax/_src/linear_algebra.py

@@ -276,6 +276,7 @@ def _iter_body(state):
         _iter_condition, _iter_body, init_state
     )
     error = jnp.max(jnp.abs(mat_m - identity))
+    # pyrefly: ignore [missing-attribute]
     is_converged = jnp.asarray(convergence, old_mat_h.dtype)  # pytype: disable=attribute-error  # lax-types # noqa: E501
     resultant_mat_h = is_converged * mat_h + (1 - is_converged) * old_mat_h
     # pyrefly: ignore [missing-attribute]

optax/_src/linesearch.py

@@ -426,6 +426,7 @@ def body_fn(
     if verbose:
       # We print information only if the linesearch failed.
       _cond_print(
+          # pyrefly: ignore [unsupported-operation]
           search_state.decrease_error > atol,
           "INFO: optax.scale_by_backtracking_linesearch:\n"
           "Backtracking linesearch failed to find a stepsize ensuring sufficent"

optax/_src/sharding_test.py

@@ -15,6 +15,7 @@
 
 """Module for testing sharding and related behavior of the optax public API."""
 
+import math
 import os
 
 from absl.testing import absltest
@@ -142,6 +143,47 @@ def test_microbatch_with_explicit_sharding(self, spec):
       )
       test_utils.assert_trees_all_equal(actual, expected)
 
+  @parameterized.named_parameters(
+      ('axis_neg1', -1, (2, 16), jax.sharding.PartitionSpec(None, 'x')),
+      (
+          'axis_neg2',
+          -2,
+          (2, 16, 4),
+          jax.sharding.PartitionSpec(None, 'x', None),
+      ),
+  )
+  def test_microbatch_negative_axis_with_explicit_sharding(
+      self, in_axes, shape, spec
+  ):
+    if utils.parse_version(jax.__version__) < utils.parse_version('0.8.1'):
+      self.skipTest('Skipping sharding-in-types test.')
+    mesh = jax.make_mesh(
+        (8,), ('x',), axis_types=(jax.sharding.AxisType.Explicit,)
+    )
+    with jax.set_mesh(mesh):
+      sharding = jax.sharding.NamedSharding(mesh, spec)
+      fun = lambda x: jnp.sum(x, axis=in_axes)
+      data = jax.device_put(
+          jnp.arange(math.prod(shape), dtype=jnp.float32).reshape(shape),
+          sharding,
+      )
+
+      microbatched_fun = optax.microbatching.microbatch(
+          fun,
+          argnums=0,
+          microbatch_size=8,
+          in_axes=in_axes,
+          accumulator=optax.microbatching.AccumulationType.SUM,
+      )
+
+      actual = microbatched_fun(data)
+      expected = fun(data)
+
+      test_utils.assert_trees_all_equal(
+          jax.tree.map(jax.typeof, actual), jax.tree.map(jax.typeof, expected)
+      )
+      test_utils.assert_trees_all_equal(actual, expected)
+
 
 if __name__ == '__main__':
   absltest.main()

optax/microbatching/_microbatching.py

@@ -100,7 +100,12 @@ def reshape_batch_axis(tree: Any, microbatch_size: int, axis: int = 0) -> Any:
   """
 
   def reshape_leaf(x):
-    new_shape = x.shape[:axis] + (-1, microbatch_size) + x.shape[axis + 1 :]
+    axis_ = axis if axis >= 0 else axis + x.ndim
+    new_shape = (
+        x.shape[:axis_]
+        + (-1, microbatch_size)
+        + x.shape[axis_ + 1 :]
+    )
     if utils.parse_version(jax.__version__) < utils.parse_version('0.7.0'):
       return x.reshape(new_shape, order='F')
 
@@ -112,14 +117,16 @@ def reshape_leaf(x):
         '0.8.1'
     ), 'microbatching with explicit sharding requires jax version >= 0.8.1.'
     spec = sharding.spec
-    if len(spec) < axis:  # The batch axis is not sharded.
+    if len(spec) < axis_:  # The batch axis is not sharded.
       new_spec = spec
     else:
-      new_spec = jax.P(*spec[:axis], None, spec[axis], *spec[axis + 1 :])
+      new_spec = jax.P(
+          *spec[:axis_], None, spec[axis_], *spec[axis_ + 1 :]
+      )
     out_sharding = jax.sharding.NamedSharding(sharding.mesh, new_spec)
 
     local_shape = sharding.shard_shape(x.shape)
-    nshards = x.shape[axis] // local_shape[axis]
+    nshards = x.shape[axis_] // local_shape[axis_]
     if microbatch_size % nshards != 0:
       raise ValueError(f'{nshards=} must evenly divide {microbatch_size=}.')
 
@@ -365,9 +372,10 @@ def _take_fn(index: int, axis: int) -> Callable[[jax.Array], jax.Array]:
   """Returns a function that takes the `index`-th element along the `axis`."""
 
   def fun(x):
-    if x.shape[axis] == 0:  # jnp.take doesn't work with zero axis size.
-      return jnp.empty_like(x, shape=x.shape[:axis] + x.shape[axis + 1 :])
-    return jnp.take(x, indices=index, axis=axis)
+    axis_ = axis if axis >= 0 else axis + x.ndim - 1
+    if x.shape[axis_] == 0:  # jnp.take doesn't work with zero axis size.
+      return jnp.empty_like(x, shape=x.shape[:axis_] + x.shape[axis_ + 1 :])
+    return jnp.take(x, indices=index, axis=axis_)
 
   return fun
 

optax/microbatching/_microbatching_test.py

@@ -200,6 +200,36 @@ def test_in_axes_invariant(self, acc):
     )(arg_axis1, arg_axis1)
     test_utils.assert_trees_all_close(result0, result1, atol=1e-6, rtol=1e-6)
 
+  @parameterized.parameters(
+      ((3, 4), -1, 2),
+      ((2, 3, 4), -1, 2),
+      ((2, 6, 4), -2, 2),
+      ((2, 3, 8, 5), -2, 4),
+      ((8, 3, 4), -3, 2),
+  )
+  def test_negative_in_axis(self, shape, in_axes, microbatch_size):
+    x = jnp.arange(np.prod(shape)).reshape(shape).astype(jnp.float32)
+    pos_axis = in_axes + x.ndim
+    fun = functools.partial(jnp.sum, axis=in_axes)
+
+    result_neg = microbatching.microbatch(
+        fun,
+        argnums=0,
+        microbatch_size=microbatch_size,
+        in_axes=in_axes,
+        accumulator=microbatching.AccumulationType.SUM,
+    )(x)
+    result_pos = microbatching.microbatch(
+        fun,
+        argnums=0,
+        microbatch_size=microbatch_size,
+        in_axes=pos_axis,
+        accumulator=microbatching.AccumulationType.SUM,
+    )(x)
+
+    test_utils.assert_trees_all_equal(result_neg, fun(x))
+    test_utils.assert_trees_all_equal(result_neg, result_pos)
+
   @parameterized.parameters(
       microbatching.AccumulationType.SUM,
       microbatching.AccumulationType.MEAN,

optax/perturbations/_make_pert.py

@@ -161,6 +161,7 @@ def stoch_estimator(
       baseline = None
 
     out = jax.vmap(stoch_estimator, in_axes=(0, None, None), out_axes=0)(
+        # pyrefly: ignore [bad-argument-type]
         jax.random.split(key, num_samples), x, baseline
     )
     return jax.tree.map(lambda x: jnp.mean(x, axis=0), out)

optax/projections/_projections.py

@@ -154,6 +154,7 @@ def projection_simplex(tree: Any, scale: jax.typing.ArrayLike = 1) -> Any:
   """
   values, unravel_fn = flatten_util.ravel_pytree(tree)
   new_values = scale * _projection_unit_simplex(values / scale)
+  # pyrefly: ignore [bad-argument-type]
   return unravel_fn(new_values)