Repository: google/jax

Issue:
forward-mode differentiation of lax.scan suddenly not working
When running a piece of code that was working earlier today, I am suddenly getting the following error
`
NotImplementedError: Forward-mode differentiation rule for 'while' not implemented
`

If necessary I can try to create a minimal repo, but I will say that if I replace a lax.scan call with a for loop then this error goes away. 



Discussion:
Sounds like it could be a regression! Though we have [tests](https://github.com/google/jax/blob/e087915c55f2716ba8f5464cf9f9f55f4f8ce93b/tests/lax_control_flow_test.py#L602) [for](https://github.com/google/jax/blob/e087915c55f2716ba8f5464cf9f9f55f4f8ce93b/tests/lax_control_flow_test.py#L635) [this](https://github.com/google/jax/blob/e087915c55f2716ba8f5464cf9f9f55f4f8ce93b/tests/lax_control_flow_test.py#L730), they might leave some case uncovered.

A minimal repro would be really helpful, since our current tests pass.
Ok, I have (finally) reduced this to a minimal repro, see here: https://pastebin.com/WzYk6Xj1

Let me know if anything doesn't make sense.

Thanks!
Thanks again for raising this! I only just started looking at it.

I think the issue (or at least one issue) was uncovered by #1175 or #1269: the loop being complained about is in the PRNG hash function. We're not dropping out of the autodiff trace properly somewhere on integer values (that is my guess at the root cause), and that's hitting the PRNG code which has while loops in it.

Some evidence for this guess is I replaced the line `key, split = random.split(key)` with `key, split = key, key` and the error went away.

Still looking...
I think I introduced the real issue in #1224, because in [these lines](https://github.com/google/jax/blob/22afc8cea1b9b5a0dca0cdfe8a0ce51df94cf9e0/jax/interpreters/ad.py#L488-L489) I'm effectively creating JVPTracers (with zeros) for integer-valued arguments. I think Dougal had some clever way of avoiding that before, and I clobbered it!

Resolving patch:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -61,7 +61,9 @@ def jvp_subtrace(master, primals, tangents):
   for x in list(primals) + list(tangents):
     if isinstance(x, Tracer):
       assert x.trace.level < trace.level
-  ans = yield map(partial(JVPTracer, trace), primals, tangents), {}
+  in_tracers = [JVPTracer(trace, x, t) if t is not zero else x
+                for x, t in zip(primals, tangents)]
+  ans = yield in_tracers, {}
   out_tracers = map(trace.full_raise, ans)
   yield unzip2([(out_tracer.primal, out_tracer.tangent)
                 for out_tracer in out_tracers])


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
@@ -972,6 +972,15 @@ def testWhileCondConstant(self):
     out = lax.while_loop(lambda _: False, lambda _: (), ())  # doesn't crash
     self.assertEqual(out, ())
 
+  def testIssue1316(self):
+    def f(carry, _):
+      c, key = carry
+      key, _ = random.split(key)
+      return (c, key), ()
+
+    key = random.PRNGKey(0)
+    api.grad(lambda c: lax.scan(f, (c, key), onp.ones(3))[0][0])(0.)  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()