Repository: google/jax
Issue: TypeError: <class 'jax.ad_util.Zero'> is not a valid Jax type when I combine custom_jvp and lax.scan

Hi all

I have a difficult bug which I don't understand. It started appearing when I defined my custom gradients with `custom_jvp` (in forward mode), before, I did it in reverse mode with `custom_gradient`.
I'm up-to-date on the master branch, as of this writing.

`TypeError: <class 'jax.ad_util.Zero'> is not a valid Jax type`

I've tried to make a simple code snippet to reproduce it, but it's still rather complex because the bug, as far as I can tell, only appears when I combine different things.
I've tried lots of combinations, and so far I've seen the bug appear only when:
1) I define the gradient with `custom_jvp`
2) The custom gradient depends on `dA` (`A` does not depend on `theta`, I think it's got something to do with this)
3) I use `lax.scan` (the bug disappears when I use a python loop)
4) `theta` is used as below, directly in `step` (this is not breaking the 'pure function' requirement, right?)

```python
import jax
import jax.numpy as np

@jax.custom_jvp
def f(A, b):
    return A @ b

def f_jvp(primals, tangents):
    A, b = primals
    dA, db = tangents
    z = f(A, b)
    dz = dA @ db
    return z, dz

f.defjvp(f_jvp)


def experiment(theta):
    def step(q, _):
        z = f(np.eye(3), np.ones(3) * theta)
        q += z[0]
        return q, q
    
    q = 0.
    q, _ = jax.lax.scan(step, q, None, 4)
    return q


experiment_grad = jax.grad(experiment)

g = experiment_grad(1.)
print(g)
```

Thanks for the help!

Rembert

Changed files:
jax/custom_derivatives.py, tests/api_test.py

diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
index dc1b40e6f..8988e30a0 100644
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -294,7 +294,8 @@ custom_jvp_call_jaxpr_p.def_abstract_eval(_custom_jvp_call_jaxpr_abstract_eval)
 
 def _custom_jvp_call_jaxpr_jvp(primals, tangents, *, fun_jaxpr, jvp_jaxpr_thunk):
   jvp_jaxpr = jvp_jaxpr_thunk()
-  outs = core.jaxpr_as_fun(jvp_jaxpr)(*(primals + tangents))
+  tangents = map(ad.instantiate_zeros, primals, tangents)
+  outs = core.jaxpr_as_fun(jvp_jaxpr)(*primals, *tangents)
   return split_list(outs, [len(outs) // 2])
 ad.primitive_jvps[custom_jvp_call_jaxpr_p] = _custom_jvp_call_jaxpr_jvp
 
diff --git a/tests/api_test.py b/tests/api_test.py
index 51742666f..f61b00da8 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2446,6 +2446,34 @@ class CustomJVPTest(jtu.JaxTestCase):
     api.eval_shape(expit, np.ones((2, 3)))
     api.eval_shape(api.grad(lambda x: expit(x).sum()), np.ones((2, 3)))
 
+  def test_jaxpr_zeros(self):
+    # from https://github.com/google/jax/issues/2657
+    @api.custom_jvp
+    def f(A, b):
+        return A @ b
+
+    def f_jvp(primals, tangents):
+        A, b = primals
+        dA, db = tangents
+        z = f(A, b)
+        dz = A @ db + dA @ b
+        return z, dz
+
+    f.defjvp(f_jvp)
+
+    def experiment(theta):
+        def step(q, _):
+            z = f(np.eye(3), np.ones(3) * theta)
+            q += z[0]
+            return q, q
+
+        q = 0.
+        q, _ = lax.scan(step, q, None, 4)
+        return q
+
+    grad(experiment)(1.)  # doesn't crash
+
+
 class CustomVJPTest(jtu.JaxTestCase):
 
   def test_basic(self):