Repository: google/jax

Issue:
Batching rules for pad and concatenate primitives not implemented
I'm interested in using JAX to compute Jacobians for functions which involve iteratively applying operations to generate a sequence. When using Autograd for this, in order to avoid indexed assignment I would create a list which is iteratively populated with the sequence values and then create an array from the list using `np.array` or `np.stack`. Attempting the same in JAX (built from source with fc4afb409b) raises a `NotImplementedError` when trying to compute the Jacobian of such a function with either `jacrev` or `jacfwd` as batching rules are not implemented for the `pad` and `concatenate` primitives respectively. 

As a minimal example

```Python
import jax.numpy as np
from jax import jacrev, jacfwd

def func(xs): 
    return np.array([x for x in xs])

jacrev_func = jacrev(func)
jacfwd_func = jacfwd(func)

xs = np.ones((5, 1))

jacrev_func(xs)
# raises NotImplementedError: Batching rule for 'pad' not implemented

jacfwd_func(xs)
# raises NotImplementedError: Batching rule for 'concatenate' not implemented
```

The same errors are raised when replacing `np.array` in the defnition of `func` with `np.stack`, `np.hstack`, `np.vstack` or `np.concatenate`.


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -64,7 +64,7 @@ def jacfwd(fun, x):
   fun = lu.wrap_init(fun)
   pushfwd = partial(jvp, fun, (x,))
   std_basis = onp.eye(onp.size(x)).reshape((-1,) + onp.shape(x)),
-  y, jac_flat = vmap(pushfwd, std_basis, out_axes=(None, 0))
+  y, jac_flat = vmap(pushfwd, out_axes=(None, -1))(std_basis)
   return jac_flat.reshape(onp.shape(y) + onp.shape(x))
 
 @curry
@@ -72,7 +72,7 @@ def jacrev(fun, x):
   fun = lu.wrap_init(fun)
   y, pullback = vjp(fun, x)
   std_basis = onp.eye(onp.size(y)).reshape((-1,) + onp.shape(y))
-  jac_flat, = vmap(pullback, std_basis, out_axes=onp.ndim(y))
+  jac_flat, = vmap(pullback, out_axes=0)(std_basis)
   return jac_flat.reshape(onp.shape(y) + onp.shape(x))
 
 def hessian(fun):
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -281,6 +281,7 @@ def moveaxis(sz, dst, src, x):
     else:
       return pack(map(partial(moveaxis, sz, dst, src), x))
   elif isinstance(aval, ShapedArray):
+    dst = (dst % aval.ndim) if dst is not None and aval.ndim else dst
     if src == dst:
       return x
     else:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1362,11 +1362,20 @@ def concatenate_transpose_rule(t, *operands, **kwargs):
     return [slice(t, start, limit) if o is None else None
             for o, start, limit in zip(operands, starts, limits)]
 
+def concatenate_batch_rule(batched_args, batch_dims, dimension, operand_shapes):
+  size = next(op.shape[bdim] for op, bdim in zip(batched_args, batch_dims)
+              if bdim is not None)
+  operands = [batching.move_dim_to_front(op, bdim) if bdim is not None
+              else broadcast(op, (size,))
+              for op, bdim in zip(batched_args, batch_dims)]
+  return concatenate(operands, dimension + 1), 0
+
 concatenate_p = standard_primitive(
     concatenate_shape_rule, concatenate_dtype_rule, 'concatenate',
     concatenate_translation_rule)
 ad.deflinear(concatenate_p, concatenate_transpose_rule)
 ad.primitive_transposes[concatenate_p] = concatenate_transpose_rule
+batching.primitive_batchers[concatenate_p] = concatenate_batch_rule
 
 
 def pad_shape_rule(operand, padding_value, padding_config):
@@ -1398,9 +1407,21 @@ def pad_transpose(t, operand, padding_value, padding_config):
 
   return [t_operand, t_padv]
 
+def pad_batch_rule(batched_args, batch_dims, padding_config):
+  operand, padding_value = batched_args
+  operand_bdim, padding_value_bdim = batch_dims
+  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
+
 pad_p = standard_primitive(pad_shape_rule, _input_dtype, 'pad')
 ad.deflinear(pad_p, pad_transpose)
 ad.primitive_transposes[pad_p] = pad_transpose
+batching.primitive_batchers[pad_p] = pad_batch_rule
 
 
 def reshape_shape_rule(operand, new_sizes, dimensions, **unused_kwargs):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -24,7 +24,7 @@
 
 import jax.numpy as np
 from jax.config import config
-from jax import jit, grad, device_get, device_put
+from jax import jit, grad, device_get, device_put, jacfwd, jacrev
 from jax.core import Primitive
 from jax.interpreters.partial_eval import def_abstract_eval
 from jax.interpreters.ad import defjvp
@@ -235,6 +235,18 @@ def test_device_put_and_get(self):
     assert isinstance(y2[1][1], onp.ndarray)
     assert onp.all(y2[1][1] == 3 * x)
 
+  def test_jacobian(self):
+    R = onp.random.RandomState(0).randn
+    A = R(4, 3)
+    x = R(3)
+
+    f = lambda x: np.dot(A, x)
+    assert onp.allclose(jacfwd(f)(x), A)
+    assert onp.allclose(jacrev(f)(x), A)
+
+    f = lambda x: np.tanh(np.dot(A, x))
+    assert onp.allclose(jacfwd(f)(x), jacrev(f)(x))
+
 
 if __name__ == '__main__':
   config.config_with_absl()
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -24,7 +24,7 @@
 from jax import test_util as jtu
 from jax.abstract_arrays import ShapedArray
 from jax import lax
-from jax.api import jit, grad, jvp, vjp, trace_to_jaxpr
+from jax.api import jit, grad, jvp, vjp, trace_to_jaxpr, jacfwd, jacrev
 from jax.api import vmap
 from jax.config import config
 from jax.core import unit
@@ -239,6 +239,49 @@ def vecvec(a, b):
     # TODO(mattjj): this fails due to an xla error in dot_general
     # assert vecvec(np.zeros((4, 2, 3)), np.zeros((3,))).shape == (4, 2)
 
+  def testPad(self):
+    R = onp.random.RandomState(0).randn
+
+    fun = lambda x: lax.pad(x, onp.float32(0), [(1, 2, 1)])
+    x = R(5, 10).astype(onp.float32)
+    ans = vmap(fun)(x)
+    expected_ans = np.stack(list(map(fun, x)))
+    self.assertAllClose(ans, expected_ans, check_dtypes=False)
+
+
+    fun = lambda x: lax.pad(x, onp.float32(0), [(1, 2, 1), (0, 1, 0)])
+    x = R(5, 10, 3).astype(onp.float32)
+    ans = vmap(fun)(x)
+    expected_ans = np.stack(list(map(fun, x)))
+    self.assertAllClose(ans, expected_ans, check_dtypes=False)
+
+  def testConcatenate(self):
+    R = lambda *shape: onp.random.RandomState(0).randn(*shape).astype(onp.float32)
+
+    fun = lambda *args: lax.concatenate(args, dimension=0)
+    x, y, z = R(10, 2, 3), R(1, 10, 3), R(4, 3)
+    ans = vmap(fun, in_axes=(0, 1, None))(x, y, z)
+    expected_ans = onp.concatenate([x, onp.swapaxes(y, 0, 1),
+                                    onp.broadcast_to(z, (10, 4, 3))], 1)
+    self.assertAllClose(ans, expected_ans, check_dtypes=False)
+
+    fun = lambda *args: lax.concatenate(args, dimension=1)
+    x, y, z = R(10, 2, 1), R(2, 3), R(2, 4, 10)
+    ans = vmap(fun, in_axes=(0, None, 2))(x, y, z)
+    expected_ans = onp.concatenate([x, onp.broadcast_to(y, (10, 2, 3)),
+                                    onp.moveaxis(z, 2, 0)], 2)
+    self.assertAllClose(ans, expected_ans, check_dtypes=False)
+
+  def testJacobianIssue54(self):
+    # test modeling the code in https://github.com/google/jax/issues/54
+
+    def func(xs):
+      return np.array([x for x in xs])
+
+    xs = np.ones((5, 1))
+    jacrev(func)(xs)  # don't crash
+    jacfwd(func)(xs)  # don't crash
+
 
 if __name__ == '__main__':
   config.config_with_absl()