Repository: google/jax

Issue:
incomplete implementation for batching rule for pad
For example:

```
def fun(x, pad_val):  # x: f32[3], pad_val: f32[]
  return lax.pad(x, pad_val, [(1, 1, 0)])

res = fun(np.ones((3,)), 7.)
# Pad the first row with 7 and the second with 8
res_batch = jax.vmap(fun)(np.ones((2, 3)), np.array([7, 8.]))
```

Raises NotImplementedError in the _pad_batch_rule:
```
  File "/Users/necula/Source/jax/jax/experimental/jax2tf/tests/batch_poly_test.py", line 119, in fun
    return lax.pad(x, pad_val, [(1, 1, 0)])
  File "/Users/necula/Source/jax/jax/_src/lax/lax.py", line 750, in pad
    return pad_p.bind(operand, padding_value, padding_config=tuple(padding_config))
  File "/Users/necula/Source/jax/jax/core.py", line 271, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/Users/necula/Source/jax/jax/interpreters/batching.py", line 149, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
  File "/Users/necula/Source/jax/jax/_src/lax/lax.py", line 3470, in _pad_batch_rule
    raise NotImplementedError  # loop and stack
```


Discussion:
I remember seeing a CL from George about this. Should this be closed?
Does the test pass?
I filed the bug so that I don't forget about it. I am happy to work on it, but I do not have the code written.
Oh i was thinking your XlaPad shape CL...Feel free to unassign this.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/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,


Test patch:
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
--- a/tests/lax_vmap_test.py
+++ b/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(