Repository: google/jax

Issue:
jit of vmapped scan results in fixed point error
``` python
import time
import jax
import jax.numpy as jnp

def fn(a):
    return jnp.cos(a)

def loop(val):
    iterations = 10
    def apply_carry(x, i):
        return jax.grad(fn, argnums=(0,))(x)[0], i

    final_val, _ = jax.lax.scan(
        apply_carry,
        val,
        jnp.arange(iterations)
    )
    return final_val

if __name__ == "__main__":
    arg = 0.5
    print(loop(arg))
    print(jax.grad(loop, argnums=(0,))(arg)) # okay
    print(jax.jacrev(loop, argnums=(0,))(arg)) # okay
    print(jax.jacfwd(loop, argnums=(0,))(arg)) # okay
    print(jax.jit(jax.jacfwd(loop, argnums=(0,)))(arg)) # error: 
```
```
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 584, in scan
    forward=True, length=length, jaxpr=jaxpr)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 890, in scan_bind
    forward=forward, length=length, jaxpr=jaxpr)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/core.py", line 120, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/interpreters/ad.py", line 252, in process_primitive
    primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 641, in _scan_jvp
    forward=forward, length=length, jaxpr=jaxpr_jvp)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 890, in scan_bind
    forward=forward, length=length, jaxpr=jaxpr)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/core.py", line 120, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/interpreters/batching.py", line 123, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 874, in _scan_batching_rule
    forward=forward, length=length, jaxpr=jaxpr_batched)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 890, in scan_bind
    forward=forward, length=length, jaxpr=jaxpr)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/core.py", line 120, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 93, in process_primitive
    return partial_eval(self, *tracers, **params)
  File "/home/yutong/venv/lib/python3.6/site-packages/jax/lax/lax_control_flow.py", line 683, in _scan_partial_eval
    raise FixedPointError
jax.lax.lax_control_flow.FixedPointError

```




Discussion:
These are fantastic test cases :) Thanks for writing these, and helping us to find and squash bugs!
Thanks - bugs are best squashed when fresh

Resolving patch:
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
@@ -19,6 +19,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import operator
+
 import numpy as onp
 
 from jax import api
@@ -619,7 +621,7 @@ def _scan_jvp(primals, tangents, forward, length, jaxpr):
     jaxpr_jvp, nonzeros_out = ad.jvp_jaxpr(jaxpr, nonzeros,
                                            instantiate=(carry_nonzeros, False))
     carry_nonzeros_out, ys_nonzeros = nonzeros_out
-    if carry_nonzeros_out == carry_nonzeros:
+    if _binary_lattice_eq(carry_nonzeros_out, carry_nonzeros):
       break
     else:
       carry_nonzeros = _binary_lattice_join(carry_nonzeros_out, carry_nonzeros)
@@ -647,19 +649,23 @@ def _scan_jvp(primals, tangents, forward, length, jaxpr):
   carry_out_dot = ad.put_zeros(ad.TangentTuple, carry_nonzeros_out, carry_out_dot)
   return core.pack((carry_out, ys)), ad.TangentTuple((carry_out_dot, ys_dot))
 
-def _binary_lattice_join(a, b):
+def _binary_lattice_fold(f, pack, a, b):
+  recur = partial(_binary_lattice_fold, f, pack)
   t = (type(a), type(b))
   if t == (tuple, tuple):
-    return tuple(map(_binary_lattice_join, a, b))
+    return pack(map(recur, a, b))
   elif t == (tuple, bool):
-    return tuple(map(_binary_lattice_join, a, (b,) * len(a)))
+    return pack(map(recur, a, (b,) * len(a)))
   elif t == (bool, tuple):
-    return tuple(map(_binary_lattice_join, (a,) * len(b), b))
+    return pack(map(recur, (a,) * len(b), b))
   elif t == (bool, bool):
-    return a or b
+    return f(a, b)
   else:
     raise TypeError((type(a), type(b)))
 
+_binary_lattice_join = partial(_binary_lattice_fold, operator.or_, tuple)
+_binary_lattice_eq = partial(_binary_lattice_fold, operator.eq, all)
+
 
 def _scan_partial_eval(trace, *tracers, **kwargs):
   jaxpr = kwargs.pop('jaxpr')
@@ -675,7 +681,7 @@ def _scan_partial_eval(trace, *tracers, **kwargs):
     jaxpr_1, jaxpr_2, sc_out = pe.partial_eval_jaxpr(jaxpr, second_components,
                                                      instantiate=(sc_carry, False))
     sc_carry_out, sc_ys = sc_out
-    if sc_carry_out == sc_carry:
+    if _binary_lattice_eq(sc_carry_out, sc_carry):
       break
     else:
       sc_carry = _binary_lattice_join(sc_carry, sc_carry_out)
@@ -858,7 +864,7 @@ def _scan_batching_rule(batched_args, batch_dims, forward, length, jaxpr):
     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:
+    if _binary_lattice_eq(carry_batched_out, carry_batched):
       break
     else:
       carry_batched = _binary_lattice_join(carry_batched_out, carry_batched)


Test patch:
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
@@ -787,6 +787,26 @@ def f(c, a):
     ans = api.vmap(lambda c, as_:            lax.scan(f, c, as_), in_axes)(c, as_)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue757(self):
+    # code from https://github.com/google/jax/issues/757
+    def fn(a):
+        return np.cos(a)
+
+    def loop(val):
+        iterations = 10
+        def apply_carry(x, i):
+            return api.grad(fn, argnums=(0,))(x)[0], i
+
+        final_val, _ = lax.scan(
+            apply_carry,
+            val,
+            np.arange(iterations)
+        )
+        return final_val
+
+    arg = 0.5
+    print(api.jit(api.jacfwd(loop, argnums=(0,)))(arg))
+
 
 if __name__ == '__main__':
   absltest.main()