Repository: google/jax

Issue:
Abstract shape computed incorrectly for `lax.scan`
The abstract eval for scan ([`_scan_abstract_eval`](https://github.com/google/jax/blob/master/jax/lax.py#L1081-L1082)) returns `aval_out`, but in the case of scan, the `aval_out` should have an additional dimension corresponding to the dimension the scan is looping over (`scan` aggregates over this axis whereas `while_loop` and `fori_loop` don't). 

This crops up when using `jit` on `lax.scan` functions that use the output of `lax.scan` (if you just return the output of `lax.scan` it works fine). Here's how to reproduce an error:

```python
import jax.numpy as np
from jax import lax, jit

def f1(bs):
    result = lax.scan(lambda a, _: a, np.ones(5), bs)
    return result
def f2(bs):
    result = lax.scan(lambda a, _: a, np.ones(5), bs)
    return result.transpose((1, 0)).transpose((1, 0))
fast_f1 = jit(f1)
fast_f2 = jit(f2)

f1(np.ones(10)) #works
f2(np.ones(10)) #works
fast_f1(np.ones(10)) #works
fast_f2(np.ones(10)) #errors 
```
In theory all of these should be the same, but because we call `transpose` and the abstract result only has one dimension, it errors when it is traced.

Any suggestions on the best way to fix? I came up with a hacky fix where I change
```python
def _scan_abstract_eval(a, bs, consts, aval_out, jaxpr):
  return maybe_tracer_tuple_to_abstract_tuple(aval_out)
```
to
```python
def _scan_abstract_eval(a, bs, consts, aval_out, jaxpr):
  aval_out = core.AbstractTuple(ShapedArray(shape=(bs[0].shape[0],) + a.shape, dtype=a.dtype) for a in aval_out)
  return maybe_tracer_tuple_to_abstract_tuple(aval_out)
```


Discussion:
Thanks for raising this!

Actually we should probably delete the `lax.scan` on master. The real one is coming in the differentiable-scan branch. Still a work-in-progress, but we hope it will be ready in the next week or two.

Unfortunately that means for now there really is no scan...

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1000,94 +1000,6 @@ def _revise_cond_jaxpr(new_pval, old_pval, jaxpr, consts):
     return new_jaxpr, new_consts
 
 
-def scan(f, a, bs):
-  """Scans over the leading axis of an array.
-
-  Arguments:
-    f: function with signature `a -> b -> a`
-    a: `a` value, or a pytree of `a` values.
-    bs: an array of `b` values, or a pytree of arrays of `b` values with the
-      same leading axis size.
-
-  Returns:
-    An array of `a` values, or a pytree of arrays of `a` values, representing
-    the result of scanning the function `f` over the leading axis of `bs`, with
-    each application producing an `a` for the next and collecting the results.
-  """
-  a, a_tree = pytree_to_flatjaxtuple(a)
-  bs, b_tree = pytree_to_flatjaxtuple(bs)  # b_tree is the same as bs_tree
-  f, out_tree = pytree_fun_to_flatjaxtuple_fun(lu.wrap_init(f), (a_tree, b_tree))
-
-  if not bs:
-    raise TypeError("bs argument to scan does not contain any arrays")
-  if any([b.ndim == 0 for b in bs]):
-    msg = "bs argument arrays must be rank >=1, got shapes {}."
-    raise TypeError(msg.format(", ".format(str(b.shape) for b in bs)))
-  if len({b.shape[0] for b in bs}) != 1:
-    msg = "arrays in bs must have equal most-major dimensions, got shapes {}."
-    raise TypeError(msg.format(", ".format(str(b.shape) for b in bs)))
-
-  a_pval = a_aval, _ = _abstractify(a)
-  bs_aval, _ = _abstractify(bs)
-  b_aval = core.AbstractTuple([ShapedArray(b.shape[1:], b.dtype) for b in bs_aval])
-  b_pval = pe.PartialVal((b_aval, core.unit))
-  jaxpr, pval_out, consts = pe.trace_to_jaxpr(f, (a_pval, b_pval))
-  aval_out, _ = pval_out
-
-  if a_tree != out_tree():
-    msg = "scanned function input and output must have identical structure"
-    raise TypeError(msg)
-  if a_aval != aval_out:
-    msg = "output shape mismatch for scanned function: {} vs {}"
-    raise TypeError(msg.format(a_aval, aval_out))
-
-  out = scan_p.bind(a, bs, core.pack(consts), aval_out=aval_out, jaxpr=jaxpr)
-  return tree_unflatten(out_tree(), out)
-
-def _scan_impl(a, bs, consts, aval_out, jaxpr):
-  length = tuple(bs)[0].shape[0]
-  state = [full((length,) + elt.shape, 0, _dtype(elt)) for elt in a]
-
-  def body_fun(i, vals):
-    a, state = vals
-    assert len(a) == len(state)
-    b = [dynamic_index_in_dim(b, i, keepdims=False) for b in bs]
-    a_out = core.eval_jaxpr(jaxpr, consts, (), a, core.pack(b))
-    state_out = [dynamic_update_index_in_dim(s, a[None, ...], i, axis=0)
-                 for a, s in zip(a_out, state)]
-    return a_out, state_out
-
-  _, out = fori_loop(0, length, body_fun, (a, state))
-  return core.pack(out)
-
-# TODO(mattjj, phawkins): figure out what to do with consts_tangents, and the
-# jaxtuple packing issues
-def _scan_jvp(primals, tangents, aval_out, jaxpr):
-  a, bs, consts_primals = primals
-  a_dot, bs_dot, consts_tangents = tangents
-
-  primal_out = scan_p.bind(a, bs, consts_primals,
-                           aval_out=aval_out, jaxpr=jaxpr)
-
-  def f_jvp(a_pt, b_pt):
-    a, a_dot = a_pt
-    b, b_dot = b_pt
-    f = lambda a, b, c: core.eval_jaxpr(jaxpr, c, (), a, b)
-    return api.jvp(f, (a, b, consts), (b, b_dot, consts_tangents))
-  tangent_out = scan(f_jvp, (a, a_dot), (b, b_dot))
-
-  return primal_out, tangent_out
-
-def _scan_abstract_eval(a, bs, consts, aval_out, jaxpr):
-  return maybe_tracer_tuple_to_abstract_tuple(aval_out)
-
-scan_p = core.Primitive("scan")
-scan_p.def_impl(_scan_impl)
-scan_p.def_abstract_eval(_scan_abstract_eval)
-xla.translations[scan_p] = partial(xla.lower_fun, _scan_impl)
-# ad.primitive_jvps[scan_p] = _scan_jvp  # TODO(mattjj, phawkins)
-
-
 def tie_in(x, y):
   return tie_in_p.bind(x, y)
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1562,45 +1562,6 @@ def testIssue514(self):
             (0, 0), lambda x: (x[0], 0),
             (1, 1), lambda x: x)
 
-  def testScanAdd(self):
-    def f(x, y):
-      return x + y
-
-    g = partial(lax.scan, f)
-    a = onp.array(7, onp.float32)
-    bs = onp.array([2, 4, -2, 6], onp.float32)
-    out = g(a, bs)
-    self.assertAllClose(out, onp.array([9, 13, 11, 17], onp.float32),
-                        check_dtypes=True)
-
-    # jtu.check_jvp(g, partial(api.jvp, g), (a, bs))
-
-  def testScanMul(self):
-    def f(x, y):
-      return x * y
-
-    g = partial(lax.scan, f)
-    a = onp.array(7, onp.float32)
-    bs = onp.array([2, 4, -2, 6], onp.float32)
-    out = g(a, bs)
-    self.assertAllClose(out, onp.array([14, 56, -112, -672], onp.float32),
-                        check_dtypes=True)
-
-    # jtu.check_jvp(g, partial(api.jvp, g), (a, bs))
-
-  def testScanJit(self):
-    @api.jit
-    def f(x, yz):
-      y, z = yz
-      return 5. * lax.exp(lax.sin(x) * lax.cos(y)) + z
-
-    a = onp.array(7, onp.float32)
-    bs = (onp.array([3., 1., -4., 1.], onp.float32),
-          onp.array([5., 9., -2., 6.], onp.float32))
-    ans = lax.scan(f, a, bs)
-    expected = onp.array([7.609, 17.445, 7.52596, 14.3389172], onp.float32)
-    self.assertAllClose(ans, expected, check_dtypes=True)
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_lhs_shape={}_rhs_shape={}"
        .format(jtu.format_shape_dtype_string(lhs_shape, dtype),