Upstream num_real_microbatches to micro_vmap and micro_grad and add unit tests.

optax/experimental/microbatching.py

@@ -461,6 +461,7 @@ def micro_vmap(
     accumulator: (
         Accumulator | AccumulationType | AccumulatorTree
     ) = AccumulationType.CONCAT,
+    num_real_microbatches: int | jax.Array | None = None,
 ) -> Function:
   """A generalized version of jax.vmap that supports microbatching.
 
@@ -494,6 +495,10 @@ def micro_vmap(
       batch axis is not needed and is too large to fit in memory. This
       accumulator can be any PyTree prefix of the outputs of `fun` to apply
       different reductions to different sub-trees.
+    num_real_microbatches: Optional number of microbatches that are actually
+      executed. If specified, microbatching will terminate early after this
+      many steps. Can be helpful to handle variable batch sizes without
+      recompilation.
 
   Returns:
     A new function with the same args and kwargs having an additional
@@ -524,6 +529,7 @@ def vmap_reduce_fn(*args, kwargs):
       microbatch_size=microbatch_size,
       accumulator=accumulator,
       in_axes=tuple(ax for ax in in_axes if ax is not None) + (0,),
+      num_real_microbatches=num_real_microbatches,
   )
 
   def wrapped_fn(*args, **kwargs):
@@ -566,7 +572,8 @@ def micro_grad(
         Accumulator | AccumulationType | AccumulatorTree
     ) = AccumulationType.SUM,
     transform_fn: Callable[[chex.ArrayTree], chex.ArrayTree] = lambda x: x,
-    metrics_fn: Callable[[chex.ArrayTree], chex.ArrayTree] = lambda x: None
+    metrics_fn: Callable[[chex.ArrayTree], chex.ArrayTree] = lambda x: None,
+    num_real_microbatches: int | jax.Array | None = None,
 ) -> ValueAndGradFn:
   """Create a function to compute, transform, and sum per-example gradients.
 
@@ -626,6 +633,10 @@ def micro_grad(
     metrics_fn: A function to apply to per-example gradients before
       transforming. Will be returned on a per-example basis as part of the
       auxiliary output, and therefore should be scalar or low-dimensional.
+    num_real_microbatches: Optional number of microbatches that are actually
+      executed. If specified, microbatching will terminate early after this
+      many steps. Can be helpful to handle variable batch sizes without
+      recompilation.
 
   Returns:
     A function that computes the value and gradient of `fun`, averaging the
@@ -657,7 +668,8 @@ def grad_fn(*args, **kwargs):
       grad_fn,
       in_axes=in_axes,
       accumulator=(accumulator, AccumulationType.CONCAT),
-      microbatch_size=microbatch_size
+      microbatch_size=microbatch_size,
+      num_real_microbatches=num_real_microbatches,
   )
   if keep_batch_dim:
     micro_fun = _with_extra_batch_axis(micro_fun, batch_argnums)

optax/experimental/microbatching_test.py

@@ -271,6 +271,16 @@ def fun(x, *, y):
     normal_vmap = jax.vmap(fun, in_axes=0)
     test_utils.assert_trees_all_equal(normal_vmap(x, y=y), custom_vmap(x, y=y))
 
+  def test_vmap_early_stopping(self):
+    x = jnp.ones(16)
+    custom_vmap = microbatching.micro_vmap(
+        jnp.sum,
+        microbatch_size=4,
+        num_real_microbatches=3,
+        accumulator=microbatching.AccumulationType.SUM,
+    )
+    test_utils.assert_trees_all_close(custom_vmap(x), 12.0)
+
   @parameterized.parameters([False, True])
   def test_micro_grad_basic(self, keep_batch_dim):
     # This function definition is valid whether or not features/targets has a
@@ -306,6 +316,13 @@ def mean_squared_loss(params, features, targets):
     test_utils.assert_trees_all_close(aux.metrics, expected_norms)
     self.assertIsNone(aux.aux)
 
+  def test_micro_grad_early_stopping(self):
+    grad_fn = microbatching.micro_grad(
+        lambda c, x: c * x.sum(), microbatch_size=4, num_real_microbatches=3
+    )
+    result, _ = grad_fn(1.0, jnp.ones(16))
+    test_utils.assert_trees_all_close(result, 12.0)
+
 
 if __name__ == '__main__':
   absltest.main()