From cff034747e7fe88e287782744bf4659fc437abe7 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:56:19 +0530 Subject: [PATCH 1/3] Canonicalize negative microbatch axes --- optax/microbatching/_microbatching.py | 22 +++++++++++++++------- optax/microbatching/_microbatching_test.py | 14 ++++++++++++++ 2 files changed, 29 insertions(+), 7 deletions(-) diff --git a/optax/microbatching/_microbatching.py b/optax/microbatching/_microbatching.py index ef7d82752..7fde93508 100644 --- a/optax/microbatching/_microbatching.py +++ b/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 diff --git a/optax/microbatching/_microbatching_test.py b/optax/microbatching/_microbatching_test.py index cb6390388..3678e9e96 100644 --- a/optax/microbatching/_microbatching_test.py +++ b/optax/microbatching/_microbatching_test.py @@ -200,6 +200,20 @@ 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) + def test_negative_in_axis(self): + x = jnp.arange(12).reshape(3, 4) + fun = functools.partial(jnp.sum, axis=-1) + + result = microbatching.microbatch( + fun, + argnums=0, + microbatch_size=2, + in_axes=-1, + accumulator=microbatching.AccumulationType.SUM, + )(x) + + test_utils.assert_trees_all_equal(result, fun(x)) + @parameterized.parameters( microbatching.AccumulationType.SUM, microbatching.AccumulationType.MEAN, From a8403230f9579a8e0756a5e949e247b61eb30141 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:06:03 +0530 Subject: [PATCH 2/3] Expand negative-axis microbatch tests Cover more input shapes, in_axes=-2, and explicit sharding. Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com> --- optax/_src/sharding_test.py | 42 ++++++++++++++++++++++ optax/microbatching/_microbatching_test.py | 30 ++++++++++++---- 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/optax/_src/sharding_test.py b/optax/_src/sharding_test.py index 98afce344..fee9257bb 100644 --- a/optax/_src/sharding_test.py +++ b/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() diff --git a/optax/microbatching/_microbatching_test.py b/optax/microbatching/_microbatching_test.py index 3678e9e96..faf5fe58e 100644 --- a/optax/microbatching/_microbatching_test.py +++ b/optax/microbatching/_microbatching_test.py @@ -200,19 +200,35 @@ 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) - def test_negative_in_axis(self): - x = jnp.arange(12).reshape(3, 4) - fun = functools.partial(jnp.sum, axis=-1) + @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 = microbatching.microbatch( + result_neg = microbatching.microbatch( fun, argnums=0, - microbatch_size=2, - in_axes=-1, + 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, fun(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, From 423b27803dfbbae90f03fa4cbb600189ce9ce957 Mon Sep 17 00:00:00 2001 From: Anas Khan <83116240+anxkhn@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:41:18 +0530 Subject: [PATCH 3/3] Silence pyrefly errors from stricter JAX type stubs --- optax/_src/linear_algebra.py | 1 + optax/_src/linesearch.py | 1 + optax/perturbations/_make_pert.py | 1 + optax/projections/_projections.py | 1 + 4 files changed, 4 insertions(+) diff --git a/optax/_src/linear_algebra.py b/optax/_src/linear_algebra.py index f37ce938c..70a3450f0 100644 --- a/optax/_src/linear_algebra.py +++ b/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] diff --git a/optax/_src/linesearch.py b/optax/_src/linesearch.py index a239cee68..8918a7a10 100644 --- a/optax/_src/linesearch.py +++ b/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" diff --git a/optax/perturbations/_make_pert.py b/optax/perturbations/_make_pert.py index 753e95f71..564e419db 100644 --- a/optax/perturbations/_make_pert.py +++ b/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) diff --git a/optax/projections/_projections.py b/optax/projections/_projections.py index 97209c9d8..ac1b456d3 100644 --- a/optax/projections/_projections.py +++ b/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)