Repository: google/jax

Issue:
ode is not working in jax 0.1.70
Here is a repro code, which works for previous version

```python
import jax
import jax.numpy as jnp
from jax.experimental.ode import odeint

def dz_dt(z, t, theta):
    """ Lotka–Volterra equations. """
    u = z[0]
    v = z[1]
    alpha, beta, gamma, delta = theta[0], theta[1], theta[2], theta[3]
    du_dt = (alpha - beta * v) * u
    dv_dt = (-gamma + delta * u) * v
    return jnp.stack([du_dt, dv_dt])

def f(z):
    y = odeint(dz_dt, z, jnp.arange(10.), jnp.ones(4))
    return jnp.sum(y)

jax.grad(f)(jnp.ones(2))
```

Running the above script raises the error `TypeError: Primal inputs to reverse-mode differentiation must be of float or complex type, got type int32`. I tried to trace the error but got no hint where `int` variables are created. I think the issue happens after https://github.com/google/jax/pull/3562.


Discussion:
A simpler repro code
```
def dz_dt(z, t):
    return jnp.stack([z[0], z[1]])

def f(z):
    y = odeint(dz_dt, z, jnp.arange(10.))
    return jnp.sum(y)

jax.grad(f)(jnp.ones(2))
```

It seems to me that the indices `0`, `1` cause the issue.
Ah, this is indeed because of #3562. Thanks for catching it!
Unfortunately I've got to go afk for a while, but I should be able to fix this tonight (if no one beats me to it).
As a temporary workaround, you can use this version:

```python
from jax.experimental.ode import _odeint_wrapper

def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf):
return _odeint_wrapper(func, rtol, atol, mxstep, y0, t, *args)
```
Thanks, @mattjj!

Resolving patch:
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
--- a/jax/experimental/ode.py
+++ b/jax/experimental/ode.py
@@ -29,6 +29,7 @@
 import jax
 import jax.numpy as jnp
 from jax import core
+from jax import dtypes
 from jax import lax
 from jax import ops
 from jax.util import safe_map, safe_zip, cache, split_list
@@ -50,16 +51,31 @@ def closure_convert(fun, in_tree, in_avals):
     jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
       wrapped_fun, in_pvals, instantiate=True, stage_out=False)
   out_tree = out_tree()
-  num_consts = len(consts)
 
-  def converted_fun(y, t, *consts_args):
-    consts, args = split_list(consts_args, [num_consts])
+  # We only want to closure convert for constants with respect to which we're
+  # differentiating. As a proxy for that, we hoist consts with float dtype.
+  # TODO(mattjj): revise this approach
+  is_float = lambda c: dtypes.issubdtype(dtypes.dtype(c), jnp.inexact)
+  (closure_consts, hoisted_consts), merge = partition_list(is_float, consts)
+  num_consts = len(hoisted_consts)
+
+  def converted_fun(y, t, *hconsts_args):
+    hoisted_consts, args = split_list(hconsts_args, [num_consts])
+    consts = merge(closure_consts, hoisted_consts)
     all_args, in_tree2 = tree_flatten((y, t, *args))
     assert in_tree == in_tree2
     out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
     return tree_unflatten(out_tree, out_flat)
 
-  return converted_fun, consts
+  return converted_fun, hoisted_consts
+
+def partition_list(choice, lst):
+  out = [], []
+  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]
+  def merge(l1, l2):
+    i1, i2 = iter(l1), iter(l2)
+    return [next(i2 if snd else i1) for snd in which]
+  return out, merge
 
 def abstractify(x):
   return core.raise_to_shaped(core.get_aval(x))


Test patch:
diff --git a/tests/ode_test.py b/tests/ode_test.py
--- a/tests/ode_test.py
+++ b/tests/ode_test.py
@@ -221,6 +221,19 @@ def f(k):
     with self.assertRaisesRegex(TypeError, "can't apply forward-mode.*"):
       jax.jacfwd(f)(3.)
 
+  @jtu.skip_on_devices("tpu")
+  def test_closure_nondiff(self):
+    # https://github.com/google/jax/issues/3584
+
+    def dz_dt(z, t):
+      return jnp.stack([z[0], z[1]])
+
+    def f(z):
+      y = odeint(dz_dt, z, jnp.arange(10.))
+      return jnp.sum(y)
+
+    jax.grad(f)(jnp.ones(2))  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()