Repository: google/jax

Issue:
Please enable vmap(scan())



Discussion:
The addition of scan() is great!  I anticipate its usage will be much wider if one can call vmap() over scanned functions.  It's not yet implemented.  
```NotImplementedError: Batching rule for 'scan' not implemented```

Thank you!

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -557,9 +557,6 @@ def map_transpose(primitive, params, jaxpr, consts, freevar_vals, args, ct):
   freevar_cts = tree_map(lambda x: x.sum(0), freevar_cts)
   return cts_out, freevar_cts
 
-def jaxpr_as_fun(jaxpr, consts, *args):
-  return core.eval_jaxpr(jaxpr, consts, (), *args)
-
 def get_nonzeros(tangent):
   if tangent is zero:
     return False
@@ -599,7 +596,7 @@ def jvp_jaxpr(jaxpr, nonzeros, instantiate):
   # jaxpr :: d -> a -> b -> (c1, c2)
   # avals = (d, a, b)
   # f :: d -> a -> b -> (c1, c2)
-  f = wrap_init(partial(jaxpr_as_fun, jaxpr.jaxpr, jaxpr.literals))
+  f = wrap_init(core.jaxpr_as_fun(jaxpr))
   f_jvp, out_nonzeros = f_jvp_traceable(jvp(f, instantiate=instantiate), nonzeros)
   # f_jvp :: (d, d') -> (a, a') -> (b, b') -> ((c1, c1'), (c2, c2'))
   tangent_avals = map(partial(strip_zeros, core.AbstractTuple(()), core.AbstractTuple),
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -29,9 +29,9 @@
 from ..abstract_arrays import ShapedArray, make_shaped_array, array_types, raise_to_shaped
 from ..ad_util import add_jaxvals_p, zeros_like_p, zeros_like_jaxval
 from ..linear_util import transformation, transformation_with_aux, wrap_init
-from ..tree_util import register_pytree_node
 from ..util import unzip2, partial, safe_map
 from . import xla
+from . import partial_eval as pe
 
 map = safe_map
 
@@ -52,10 +52,10 @@ def batch_transform(size, in_dims, out_dim_dst, vals):
   with new_master(BatchTrace) as master:
     trace = BatchTrace(master, core.cur_sublevel())
     in_tracers = map(partial(BatchTracer, trace), vals, in_dims)
-    out_tracer = yield in_tracers, {}
-    out_tracer = trace.full_raise(out_tracer)
+    ans = yield in_tracers, {}
+    out_tracer = trace.full_raise(ans)
     out_val, out_dim = out_tracer.val, out_tracer.batch_dim
-    del master
+    del master, out_tracer
   yield moveaxis(size, out_dim_dst, out_dim, out_val)
 
 
@@ -75,6 +75,7 @@ class BatchTracer(Tracer):
   __slots__ = ['val', 'batch_dim']
 
   def __init__(self, trace, val, batch_dim):
+    assert core.skip_checks or type(batch_dim) in (int, tuple)
     self.trace = trace
     self.val = val
     self.batch_dim = batch_dim
@@ -365,3 +366,150 @@ def handle_scalar_broadcasting(nd, x, bdim):
     return x
   else:
     return x.reshape(x.shape + (1,) * (nd - x.ndim))
+
+
+# TODO(mattjj): try to de-duplicate utility functions with above
+
+def where_batched(bdim):
+  t = type(bdim)
+  if t is tuple:
+    return tuple(map(where_batched, bdim))
+  elif t in (int, type(None)):
+    return bdim is not None
+  else:
+    raise TypeError(t)
+
+def bools_to_bdims(bdim, batched_indicator_tree):
+  t = type(batched_indicator_tree)
+  if t is tuple:
+    return tuple(map(partial(bools_to_bdims, bdim), batched_indicator_tree))
+  elif t is bool:
+    return bdim if batched_indicator_tree else None
+  else:
+    raise TypeError(t)
+
+def instantiate_bdim(size, axis, instantiate, bdim, x):
+  """Instantiate or move a batch dimension to position `axis`.
+
+  Ensures that `x` is at least as high on the batched lattice as `instantiate`.
+
+  Args:
+    size: int, size of the axis to instantiate.
+    axis: int, where to instantiate or move the batch dimension.
+    instantiate: tuple-tree of booleans, where the tree structure is a prefix of
+      the tree structure in x, indicating whether to instantiate the batch
+      dimension at the corresponding subtree in x.
+    bdim: tuple-tree of ints or NoneTypes, with identical tree structure to
+      `instantaite`, indicating where the batch dimension exists in the
+      corresponding subtree of x.
+    x: JaxType value on which to instantiate or move batch dimensions.
+
+  Returns:
+    A new version of `x` with instantiated batch dimensions.
+  """
+  def _inst(instantiate, bdim, x):
+    if type(instantiate) is tuple:
+      assert type(bdim) is tuple  # instantiate at same granularity as bdim
+      return core.pack(map(_inst, instantiate, bdim, x))
+    elif type(instantiate) is bool:
+      if not instantiate:
+        if bdim is None:
+          return x
+        elif type(bdim) is int:
+          return moveaxis2(bdim, axis, x)
+        else:
+          raise TypeError(type(bdim))
+      else:
+        if bdim is None:
+          return broadcast2(size, axis, x)
+        elif type(bdim) is int:
+          return moveaxis2(bdim, axis, x)
+        else:
+          raise TypeError(type(bdim))
+    else:
+      raise TypeError(type(instantiate))
+
+  return _inst(instantiate, bdim, x)
+
+def moveaxis2(src, dst, x):
+  if src == dst:
+    return x
+  else:
+    return _moveaxis2(src, dst, x, get_aval(x))
+
+def _moveaxis2(src, dst, x, aval):
+  if type(aval) is JaxTuple:
+    return core.pack(map(partial(_moveaxis2, src, dst), x, aval))
+  else:
+    perm = [i for i in range(onp.ndim(x)) if i != src]
+    perm.insert(dst, src)
+    return x.transpose(perm)
+
+def broadcast2(size, axis, x):
+  return _broadcast2(size, axis, x, get_aval(x))
+
+def _broadcast2(size, axis, x, aval):
+  if type(aval) is JaxTuple:
+    return core.pack(map(partial(_broadcast2, size, axis), x, aval))
+  else:
+    # see comment at the top of this section
+    if isinstance(x, onp.ndarray) or onp.isscalar(x):
+      return onp.broadcast_to(x, (size,) + onp.shape(x))
+    else:
+      return x.broadcast((size,))  # should be a JAX arraylike
+
+def _promote_aval_rank(n, batched, aval):
+  assert isinstance(aval, core.AbstractValue)
+  if isinstance(aval, core.AbstractTuple):
+    t = type(batched)
+    if t is tuple:
+      return core.AbstractTuple(map(partial(_promote_aval_rank, n), batched, aval))
+    elif t is bool:
+      if batched:
+        return core.AbstractTuple(map(partial(_promote_aval_rank, n, batched), aval))
+      else:
+        return aval
+    else:
+      raise TypeError(t)
+  else:
+    if batched:
+      return ShapedArray((n,) + aval.shape, aval.dtype)
+    else:
+      return aval
+
+def batch_jaxpr(jaxpr, size, is_batched, instantiate):
+  f = wrap_init(core.jaxpr_as_fun(jaxpr))
+  f_batched, where_out_batched = batched_traceable(f, size, is_batched, instantiate)
+  in_avals = map(partial(_promote_aval_rank, size), is_batched, jaxpr.in_avals)
+  in_pvals = [pe.PartialVal((aval, core.unit)) for aval in in_avals]
+  jaxpr_out, pval_out, literals_out = pe.trace_to_jaxpr(
+      f_batched, in_pvals, instantiate=True)
+  out_aval, _ = pval_out
+  jaxpr_out = core.TypedJaxpr(jaxpr_out, literals_out, in_avals, out_aval)
+  return jaxpr_out, where_out_batched()
+
+@transformation_with_aux
+def batched_traceable(size, is_batched, instantiate, *vals):
+  in_dims = bools_to_bdims(0, is_batched)
+  with new_master(BatchTrace) as master:
+    trace = BatchTrace(master, core.cur_sublevel())
+    in_tracers = map(partial(BatchTracer, trace), vals, in_dims)
+    ans = yield in_tracers, {}
+    out_tracer = trace.full_raise(ans)
+    out_val, out_dim = out_tracer.val, out_tracer.batch_dim
+    del master, out_tracer
+  out_val = instantiate_bdim(size, 0, instantiate, out_dim, out_val)
+  yield out_val, _binary_lattice_join(where_batched(out_dim), instantiate)
+
+def _binary_lattice_join(a, b):
+  t = (type(a), type(b))
+  if t == (tuple, tuple):
+    return tuple(map(_binary_lattice_join, a, b))
+  elif t == (tuple, bool):
+    return tuple(map(_binary_lattice_join, a, (b,) * len(a)))
+  elif t == (bool, tuple):
+    return tuple(map(_binary_lattice_join, (a,) * len(b), b))
+  elif t == (bool, bool):
+    return a or b
+  else:
+    raise TypeError((type(a), type(b)))
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -35,7 +35,7 @@
 from jax.interpreters import xla
 from jax.interpreters import ad
 from jax.util import partial, unzip2, safe_map, safe_zip
-from jax.tree_util import build_tree, tree_unflatten
+from jax.tree_util import build_tree, tree_unflatten, tree_map
 from jax import ad_util
 
 map = safe_map
@@ -236,6 +236,8 @@ def _while_loop_batching_rule(batched_args, batch_dims, cond_consts,
   # traceable Python functions using `core.eval_jaxpr`. Then we can batch them
   # using `batching.batch_transform` (the transform underlying `api.vmap`). This
   # code also avoids broadcasting `cond_tracer_consts` and `body_tracer_consts`.
+  # TODO(mattjj): Revise this using scan machinery (and fixed-point the loop
+  # carry instead of lifting it all the way!)
   init_val, cond_tracer_consts, body_tracer_consts = batched_args
   init_val_bd, cond_tracer_consts_bd, body_tracer_consts_bd = batch_dims
 
@@ -501,6 +503,8 @@ def _update_arrays(i, aval, xs, x):
   else:
     return lax.dynamic_update_index_in_dim(xs, x[None, ...], i, axis=0)
 
+class FixedPointError(Exception): pass
+
 
 def scan(f, init, xs):
   """Scan a function over leading array axes while carrying along state.
@@ -814,7 +818,6 @@ def _move_stuff_and_add_add(typed_jaxpr):
 def _add_any_eqn(tot, a, b):
   return core.JaxprEqn([a, b], [tot], ad_util.add_jaxvals_p, (), False, False, {})
 
-
 # transpose_jaxpr :: (res -> a -> b) -> (res -> CT b -> CT a)
 def _transpose_jaxpr(jaxpr):
   assert len(jaxpr.in_avals) == 2
@@ -837,7 +840,42 @@ def _make_typed_jaxpr(traceable, in_avals):
   return core.TypedJaxpr(jaxpr, consts, in_avals, out_aval)
 
 
-class FixedPointError(Exception): pass
+def _scan_batching_rule(batched_args, batch_dims, forward, length, jaxpr):
+  consts, init, xs = batched_args
+  consts_bdim, init_bdim, xs_bdim = batch_dims
+
+  sizes = lax._reduce(set.union, map(batching.dimsize, batch_dims, batched_args))
+  size = sizes.pop()
+  assert not sizes
+
+  consts_batched = batching.where_batched(consts_bdim)
+  init_batched = batching.where_batched(init_bdim)
+  xs_batched = batching.where_batched(xs_bdim)
+
+  carry_batched = init_batched
+  for _ in range(1000):
+    which_batched = (consts_batched, carry_batched, xs_batched)
+    jaxpr_batched, batched_out = batching.batch_jaxpr(jaxpr, size, which_batched,
+                                                      instantiate=(carry_batched, False))
+    carry_batched_out, ys_batched = batched_out
+    if carry_batched_out == carry_batched:
+      break
+    else:
+      carry_batched = _binary_lattice_join(carry_batched_out, carry_batched)
+  else:
+    raise FixedPointError
+
+  consts_batched = batching.instantiate_bdim(size, 0, consts_batched, consts_bdim, consts)
+  init_batched = batching.instantiate_bdim(size, 0, carry_batched, init_bdim, init)
+  xs_batched = batching.instantiate_bdim(size, 1, xs_batched, xs_bdim, xs)
+
+  carry_out, ys = scan_p.bind(
+      consts_batched, init_batched, xs_batched,
+      forward=forward, length=length, jaxpr=jaxpr_batched)
+
+  carry_out_bdim = batching.bools_to_bdims(0, carry_batched)
+  ys_bdim = batching.bools_to_bdims(1, ys_batched)
+  return core.pack((carry_out, ys)), (carry_out_bdim, ys_bdim)
 
 
 # We use a custom bind for scan just to add some error checks
@@ -858,3 +896,4 @@ def scan_bind(consts, init, xs, forward, length, jaxpr):
 ad.primitive_transposes[scan_p] = _scan_transpose
 pe.custom_partial_eval_rules[scan_p] = _scan_partial_eval
 xla.translations[scan_p] = partial(xla.lower_fun, _scan_impl)
+batching.primitive_batchers[scan_p] = _scan_batching_rule


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -889,66 +889,6 @@ def f(R):
 
     H = hessian(f)(R)  # don't crash on UnshapedArray
 
-  def testWhileLoop(self):
-    def fun(x):
-      return lax.while_loop(lambda x: x < 3, lambda x: x + 2, x)
-
-    ans = vmap(fun)(onp.array([0, 1, 2, 3]))
-    expected = onp.array([4, 3, 4, 3])
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
-    fun = jit(fun)
-    ans = vmap(fun)(onp.array([0, 1, 2, 3]))
-    expected = onp.array([4, 3, 4, 3])
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
-  def testWhileLoopCondConstsBatched(self):
-    def fun(x, y):
-      return lax.while_loop(lambda x: x < y, lambda x: x + 2, x)
-
-    ans = vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
-    expected = onp.array([2, 4])
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
-  def testWhileLoopBodyConstsBatched(self):
-    def fun(x, y):
-      return lax.while_loop(lambda x: x < 3, lambda x: x + y, x)
-
-    ans = vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
-    expected = onp.array([4, 3])
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
-  def testWhileLoopTuple(self):
-    def cond_fun(loop_carry):
-      x, y = loop_carry
-      return x + y < 5
-
-    def body_fun(loop_carry):
-      x, y = loop_carry
-      x = x + 1
-      return x, y
-
-    def fun(x, y):
-      return lax.while_loop(cond_fun, body_fun, (x, y))
-
-    ans = vmap(fun)(onp.array([0, 0]), onp.array([1, 2]))
-    expected = (onp.array([4, 3]), onp.array([1, 2]))
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
-  def testForiLoop(self):
-    def body_fun(i, loop_carry):
-      x, y = loop_carry
-      x = x + 1
-      y = y + 2
-      return x, y
-
-    def fun(x):
-      return lax.fori_loop(0, 10, body_fun, (x, 0))
-
-    ans = vmap(fun)(onp.array([0, 1]))
-    expected = (onp.array([10, 11]), onp.array([20, 20]))
-    self.assertAllClose(ans, expected, check_dtypes=False)
-
   def testIssue489(self):
     def f(key):
       def body_fn(uk):
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -18,6 +18,7 @@
 
 import collections
 from functools import partial
+import itertools
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -29,6 +30,7 @@
 from jax import core
 from jax import lax
 from jax import test_util as jtu
+from jax.util import unzip2
 import jax.numpy as np  # scan tests use numpy
 
 def scan_reference(f, init, xs):
@@ -234,6 +236,66 @@ def body_fun(state):
       self.assertAllClose(cfun(x, num), onp.sum(x[:num]), check_dtypes=False)
       self.assertAllClose(cfun(x, num), onp.sum(x[:num]), check_dtypes=False)
 
+  def testWhileLoopBatched(self):
+    def fun(x):
+      return lax.while_loop(lambda x: x < 3, lambda x: x + 2, x)
+
+    ans = api.vmap(fun)(onp.array([0, 1, 2, 3]))
+    expected = onp.array([4, 3, 4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    fun = api.jit(fun)
+    ans = api.vmap(fun)(onp.array([0, 1, 2, 3]))
+    expected = onp.array([4, 3, 4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testWhileLoopCondConstsBatched(self):
+    def fun(x, y):
+      return lax.while_loop(lambda x: x < y, lambda x: x + 2, x)
+
+    ans = api.vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
+    expected = onp.array([2, 4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testWhileLoopBodyConstsBatched(self):
+    def fun(x, y):
+      return lax.while_loop(lambda x: x < 3, lambda x: x + y, x)
+
+    ans = api.vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
+    expected = onp.array([4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testWhileLoopTupleBatched(self):
+    def cond_fun(loop_carry):
+      x, y = loop_carry
+      return x + y < 5
+
+    def body_fun(loop_carry):
+      x, y = loop_carry
+      x = x + 1
+      return x, y
+
+    def fun(x, y):
+      return lax.while_loop(cond_fun, body_fun, (x, y))
+
+    ans = api.vmap(fun)(onp.array([0, 0]), onp.array([1, 2]))
+    expected = (onp.array([4, 3]), onp.array([1, 2]))
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testForiLoopBatched(self):
+    def body_fun(i, loop_carry):
+      x, y = loop_carry
+      x = x + 1
+      y = y + 2
+      return x, y
+
+    def fun(x):
+      return lax.fori_loop(0, 10, body_fun, (x, 0))
+
+    ans = api.vmap(fun)(onp.array([0, 1]))
+    expected = (onp.array([10, 11]), onp.array([20, 20]))
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testForiLoopBasic(self):
     def count(num):
       def body_fun(i, tot):
@@ -432,12 +494,12 @@ def cond(x):
     self.assertEqual(out, (7, 10))
 
   @parameterized.named_parameters(
-      {"testcase_name": "jit_scan={}_jit_f={}".format(jit_scan, jit_f),
+      {"testcase_name": "_jit_scan={}_jit_f={}".format(jit_scan, jit_f),
        "jit_scan": jit_scan, "jit_f": jit_f}
       for jit_scan in [False, True]
       for jit_f in [False, True])
   def testScanImpl(self, jit_scan, jit_f):
-    d = np.zeros(2)
+    d = np.array([1., 2.])
     def f(c, a):
       assert a.shape == (3,)
       assert c.shape == (4,)
@@ -461,12 +523,12 @@ def f(c, a):
     self.assertAllClose(ans, expected, check_dtypes=False)
 
   @parameterized.named_parameters(
-      {"testcase_name": "jit_scan={}_jit_f={}".format(jit_scan, jit_f),
+      {"testcase_name": "_jit_scan={}_jit_f={}".format(jit_scan, jit_f),
        "jit_scan": jit_scan, "jit_f": jit_f}
       for jit_scan in [False, True]
       for jit_f in [False, True])
   def testScanJVP(self, jit_scan, jit_f):
-    d = np.zeros(2)
+    d = np.array([1., 2.])
     def f(c, a):
       assert a.shape == (3,)
       assert c.shape == (4,)
@@ -490,12 +552,12 @@ def f(c, a):
     self.assertAllClose(ans, expected, check_dtypes=False)
 
   @parameterized.named_parameters(
-      {"testcase_name": "jit_scan={}_jit_f={}".format(jit_scan, jit_f),
+      {"testcase_name": "_jit_scan={}_jit_f={}".format(jit_scan, jit_f),
        "jit_scan": jit_scan, "jit_f": jit_f}
       for jit_scan in [False, True]
       for jit_f in [False, True])
   def testScanLinearize(self, jit_scan, jit_f):
-    d = np.zeros(2)
+    d = np.array([1., 2.])
     def f(c, a):
       assert a.shape == (3,)
       assert c.shape == (4,)
@@ -519,12 +581,12 @@ def f(c, a):
     self.assertAllClose(ans, expected, check_dtypes=False)
 
   @parameterized.named_parameters(
-      {"testcase_name": "jit_scan={}_jit_f={}".format(jit_scan, jit_f),
+      {"testcase_name": "_jit_scan={}_jit_f={}".format(jit_scan, jit_f),
        "jit_scan": jit_scan, "jit_f": jit_f}
       for jit_scan in [False, True]
       for jit_f in [False, True])
   def testScanGrad(self, jit_scan, jit_f):
-    d = np.zeros(2)
+    d = np.ones(2)
     def f(c, a):
       assert a.shape == (3,)
       assert c.shape == (4,)
@@ -551,9 +613,9 @@ def testScanRnn(self):
     r = npr.RandomState(0)
 
     n_in = 4
-    n_hid = 3
-    n_out = 2
-    length = 5
+    n_hid = 2
+    n_out = 1
+    length = 3
 
     W_trans = r.randn(n_hid, n_hid + n_in)
     W_out = r.randn(n_out, n_hid + n_in)
@@ -587,11 +649,18 @@ def loss(params, inputs, targets):
     # gradient evaluation doesn't crash
     api.grad(loss)(params, inputs, targets)
 
-    # gradient is zero in the right place
-    predictions = rnn(params, inputs)
-    ans = api.grad(loss)(params, inputs, predictions)
-    expected = (onp.zeros_like(W_trans), onp.zeros_like(W_out))
-    self.assertAllClose(ans, expected, check_dtypes=False)
+    # gradient check passes
+    jtu.check_grads(loss, (params, inputs, targets), order=1)
+
+    # we can vmap to batch things
+    batch_size = 7
+    batched_inputs = r.randn(batch_size, length, n_in)
+    batched_targets = r.randn(batch_size, length, n_out)
+    batched_loss = api.vmap(lambda x, y: loss(params, x, y))
+    losses = batched_loss(batched_inputs, batched_targets)
+    expected = onp.stack(list(map(lambda x, y: loss(params, x, y),
+                                  batched_inputs, batched_targets)))
+    self.assertAllClose(losses, expected, check_dtypes=False)
 
   def testIssue711(self):
     # Tests reverse-mode differentiation through a scan for which the scanned
@@ -649,6 +718,75 @@ def f(c, a):
     jtu.check_grads(lambda c, as_: lax.scan(f, c, as_), (c, as_),
                     modes=["fwd", "rev"], order=2)
 
+  @parameterized.named_parameters(
+      {"testcase_name": "_jit_scan={}_jit_f={}_in_axes={}".format(
+          jit_scan, jit_f, in_axes),
+       "jit_scan": jit_scan, "jit_f": jit_f, "in_axes": in_axes}
+      for jit_scan in [False, True]
+      for jit_f in [False, True]
+      for in_axes in itertools.product([None, 0, 1], [None, 0, 1, 2])
+      if in_axes != (None, None))
+  def testScanVmap(self, jit_scan, jit_f, in_axes):
+    d = np.array([1., 2.])
+    def f(c, a):
+      assert a.shape == (3,)
+      assert c.shape == (4,)
+      b = np.sum(np.sin(a)) + np.sum(np.sin(c)) + np.sum(np.sin(d))
+      c = np.sin(c * b)
+      assert b.shape == ()
+      return c, b
+
+    if jit_f:
+      f = api.jit(f)
+    if jit_scan:
+      scan = api.jit(lax.scan, (0,))
+    else:
+      scan = lax.scan
+
+    as_shape = [5, 3]
+    c_shape = [4]
+
+    c_bdim, as_bdim = in_axes
+    if c_bdim is not None:
+      c_shape.insert(c_bdim, 7)
+    if as_bdim is not None:
+      as_shape.insert(as_bdim, 7)
+
+    r = onp.random.RandomState(0)
+    as_ = r.randn(*as_shape)
+    c = r.randn(*c_shape)
+
+    ans = api.vmap(lambda c, as_:                scan(f, c, as_), in_axes)(c, as_)
+    expected = api.vmap(lambda c, as_: scan_reference(f, c, as_), in_axes)(c, as_)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testScanVmapTuples(self):
+    def f(c, a):
+      a1, a2 = a
+      c1, c2 = c
+      b = np.sum(np.cos(a1)) * np.sum(np.tan(c2 * a2))
+      c = c1 * np.sin(np.sum(a1 * a2)), c2 * np.cos(np.sum(a1))
+      return c, b
+
+    in_axes = (0, (1, 2))
+
+    r = onp.random.RandomState(0)
+    as_ = (r.randn(3, 7), r.randn(3, 4, 7))
+    c = (r.randn(7, 2), r.randn(7))
+
+    expected_c_out, expected_bs = [], []
+    for i in range(7):
+      c_out, bs = lax.scan(f, (c[0][i], c[1][i]), (as_[0][:,i], as_[1][:,:,i]))
+      expected_c_out.append(c_out)
+      expected_bs.append(bs)
+    expected_c_out_0, expected_c_out_1 = unzip2(expected_c_out)
+    expected_c_out = (np.stack(expected_c_out_0), np.stack(expected_c_out_1))
+    expected_bs = np.stack(expected_bs)
+    expected = expected_c_out, expected_bs
+
+    ans = api.vmap(lambda c, as_:            lax.scan(f, c, as_), in_axes)(c, as_)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()