Merge pull request jax-ml#5762 from hawkinsp:padbatch

PiperOrigin-RevId: 358070219

Author: jax authors

docs/CHANGELOG.md

@@ -16,6 +16,7 @@ PLEASE REMEMBER TO CHANGE THE '..master' WITH AN ACTUAL TAG in GITHUB LINK.
   * Added {func}`jax.experimental.jax2tf.call_tf` to call TensorFlow functions
     from JAX ([#5627](https://github.com/google/jax/pull/5627)
     and [README](https://github.com/google/jax/blob/master/jax/experimental/jax2tf/README.md#calling-tensorflow-functions-from-jax)).
+  * Extended the batching rule for `lax.pad` to support batching of the padding values.
 * Bug fixes:
   * {func}`jax.numpy.take` properly handles negative indices
     ([#5768](https://github.com/google/jax/pull/5768))

jax/_src/lax/lax.py

@@ -3536,13 +3536,22 @@ def t_op():
 def _pad_batch_rule(batched_args, batch_dims, *, padding_config):
   operand, padding_value = batched_args
   operand_bdim, padding_value_bdim = batch_dims
+  if operand_bdim is None:
+    operand_bdim = 0
+    operand = broadcast(operand, (padding_value.shape[padding_value_bdim],))
+
+  padding_config = list(padding_config)
+  padding_config.insert(operand_bdim, (0, 0, 0))
   if padding_value_bdim is None:
-    assert operand_bdim is not None
-    padding_config = list(padding_config)
-    padding_config.insert(operand_bdim, (0, 0, 0))
     return pad(operand, padding_value, padding_config), operand_bdim
-  else:
-    raise NotImplementedError  # loop and stack
+
+  assert padding_value_bdim == 0, padding_value_bdim
+
+  x = pad(operand, _zero(operand), padding_config)
+  mask = pad(full_like(operand, True, np.bool_), False, padding_config)
+  broadcasted_padding = broadcast_in_dim(padding_value, x.shape,
+                                         (operand_bdim,))
+  return select(mask, x, broadcasted_padding), operand_bdim
 
 def _pad_translation_rule(c, operand, padding_value, *, padding_config):
   return xops.Pad(operand, padding_value,

tests/lax_vmap_test.py

@@ -369,13 +369,13 @@ def testReshape(self, arg_shape, out_shape, dtype, dimensions, bdims):
        .format(jtu.format_shape_dtype_string(shape, dtype), pads, bdims),
        "shape": shape, "dtype": dtype, "pads": pads, "bdims": bdims}
       for shape in [(2, 3)]
-      for bdims in all_bdims(shape)
+      for bdims in all_bdims(shape, ())
       for dtype in default_dtypes
       for pads in [[(1, 2, 1), (0, 1, 0)]]))
   def testPad(self, shape, dtype, pads, bdims):
     rng = jtu.rand_small(self.rng())
-    fun = lambda operand: lax.pad(operand, np.array(0, dtype), pads)
-    self._CheckBatching(fun, 5, bdims, (shape,), (dtype,), rng)
+    fun = lambda operand, padding: lax.pad(operand, padding, pads)
+    self._CheckBatching(fun, 5, bdims, (shape, ()), (dtype, dtype), rng)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_predshape={}_argshapes={}_bdims={}".format(