Repository: google/jax

Issue:
AssertionError when taking grad of odeint with outer scope variable
Hello!

An AssertionError arises when taking the grad of odeint with a variable which is outer scope.

My setup is MacBook Pro 13 inch 2019 with MacOS Catalina 10.15.2. I have compiled jax and jaxlib from source on the current master branch.

Reproduction code (x is the outer scope variable):


    from jax.experimental.ode import odeint
    from jax import jit, grad, value_and_grad, vmap
    import jax.numpy as np

    @jit
    def experiment(x):
    
        def model(y, t):
            dydt = -x * y
            return dydt

        history = odeint(model, 1., np.arange(0, 10, 0.1))
        return history[-1]
    
    experiment = value_and_grad(experiment)
    t = np.arange(0., 1., 0.01)
    h, g = vmap(experiment)(t)

Running this gives the following output:

    Traceback (most recent call last):
      File "/Users/kstorm/PycharmProjects/finger_model/venv/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code
        exec(code_obj, self.user_global_ns, self.user_ns)
      File "<ipython-input-2-a7e55fc52a04>", line 1, in <module>
        runfile('/Users/kstorm/PycharmProjects/finger_model/issue.py', wdir='/Users/kstorm/PycharmProjects/finger_model')
      File "/Users/kstorm/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/193.5233.109/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
      File "/Users/kstorm/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/193.5233.109/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
        exec(compile(contents+"\n", file, 'exec'), glob, loc)
      File "/Users/kstorm/PycharmProjects/finger_model/issue.py", line 22, in <module>
        h, g = vmap(experiment)(t)
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/api.py", line 759, in batched_fun
        lambda: _flatten_axes(out_tree(), out_axes))
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/interpreters/batching.py", line 34, in batch
        return batched_fun.call_wrapped(*in_vals)
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/linear_util.py", line 150, in call_wrapped
        ans = self.f(*args, **dict(self.params, **kwargs))
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/api.py", line 428, in value_and_grad_f
        ans, vjp_py = _vjp(f_partial, *dyn_args)
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/api.py", line 1389, in _vjp
        out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/interpreters/ad.py", line 106, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
      File "/Users/kstorm/PycharmProjects/finger_model/jax/jax/interpreters/ad.py", line 97, in linearize
        assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)
    AssertionError

I know the problem can be solved by passing x as an argument to odeint and adding an argument to 'def model' like this:

    @jit
    def experiment(x):
    
        def model(y, t, a):
            dydt = -a * y
            return dydt

        history = odeint(model, 1., np.arange(0, 10, 0.1), x)
        return history[-1]`

But I wonder why it doesn't work with outer scope variables.

Thanks in advance!


Discussion:
The need for manual closure conversion is a limitation of `jax.experimental.ode.odeint`, and a consequence of how higher-order primitives are set up in JAX's internals. This issue can also come up outside of the grad or vmap context.

It would be an enhancement to `odeint` to accept functions that close over arguments. Until then, we might consider it a documentation issue, as there ought to be more guidance around it.
@froystig Could you give some more insight into what this error means?  I'm hitting it without using `odeint`, but I don't have a good MWE to post yet.
@NeilGirdhar are you differentiating through a `while_loop`? See #2129. Otherwise you'll have to give us more hints!

@killianstorm To add on to what @froystig said: `jax.experimental.ode.odeint` is not actually a primitive (like the control flow primitives `lax.cond`, `lax.scan`, etc. which correctly handle closed-over tracers in their function-valued arguments). It's just a function with a `jax.custom_vjp` rule defined. And higher-order functions with `jax.custom_vjp` rules can't automatically handle closed-over tracers in their function-valued arguments.

We should document this constraint on `jax.custom_vjp` functions better (hence the "documentation" tag on this issue), but it's mentioned briefly in the last example (a `fixed_point` function) in the [custom_vjp/jvp tutorial](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html), specifically in the paragraph [just before this heading](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#Basic-usage-of-jax.custom_jvp-and-jax.custom_vjp-APIs).

Actually, the reason why `odeint` still lives in `jax.experimental` rather than being included directly in `lax` or something is precisely that we want to upgrade it to handle closed-over tracers automatically, like the other functions/primitives in `lax` which take function-valued arguments. We just haven't done it yet! When we do, it'll graduate out of `jax.experimental`.
@mattjj You got it!  I forgot to add my vjp.   Thanks for the lightning fast response.  I really appreciate it!
@froystig, @mattjj I understand now, thanks for the explanation!

Resolving patch:
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
--- a/jax/experimental/ode.py
+++ b/jax/experimental/ode.py
@@ -31,15 +31,39 @@
 from jax import core
 from jax import lax
 from jax import ops
-from jax.util import safe_map, safe_zip
+from jax.util import safe_map, safe_zip, cache, split_list
+from jax.api_util import flatten_fun_nokwargs
 from jax.flatten_util import ravel_pytree
-from jax.tree_util import tree_map
+from jax.tree_util import tree_map, tree_flatten, tree_unflatten
+from jax.interpreters import partial_eval as pe
 from jax import linear_util as lu
 
 map = safe_map
 zip = safe_zip
 
 
+@cache()
+def closure_convert(fun, in_tree, in_avals):
+  in_pvals = [pe.PartialVal.unknown(aval) for aval in in_avals]
+  wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
+  with core.initial_style_staging():
+    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])
+    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
+
+def abstractify(x):
+  return core.raise_to_shaped(core.get_aval(x))
+
 def ravel_first_arg(f, unravel):
   return ravel_first_arg_(lu.wrap_init(f), unravel).call_wrapped
 
@@ -159,8 +183,12 @@ def _check_arg(arg):
       msg = ("The contents of odeint *args must be arrays or scalars, but got "
              "\n{}.")
       raise TypeError(msg.format(arg))
-  tree_map(_check_arg, args)
-  return _odeint_wrapper(func, rtol, atol, mxstep, y0, t, *args)
+
+  flat_args, in_tree = tree_flatten((y0, t[0], *args))
+  in_avals = tuple(map(abstractify, flat_args))
+  converted, consts = closure_convert(func, in_tree, in_avals)
+
+  return _odeint_wrapper(converted, rtol, atol, mxstep, y0, t, *consts, *args)
 
 @partial(jax.jit, static_argnums=(0, 1, 2, 3))
 def _odeint_wrapper(func, rtol, atol, mxstep, y0, ts, *args):


Test patch:
diff --git a/tests/ode_test.py b/tests/ode_test.py
--- a/tests/ode_test.py
+++ b/tests/ode_test.py
@@ -181,6 +181,37 @@ def test_disable_jit_odeint_with_vmap(self):
       f = lambda x0: odeint(lambda x, _t: x, x0, t)
       jax.vmap(f)(x0_eval)  # doesn't crash
 
+  @jtu.skip_on_devices("tpu")
+  def test_grad_closure(self):
+    # simplification of https://github.com/google/jax/issues/2718
+    def experiment(x):
+      def model(y, t):
+        return -x * y
+      history = odeint(model, 1., np.arange(0, 10, 0.1))
+      return history[-1]
+    jtu.check_grads(experiment, (0.01,), modes=["rev"], order=1)
+
+  @jtu.skip_on_devices("tpu")
+  def test_grad_closure_with_vmap(self):
+    # https://github.com/google/jax/issues/2718
+    @jax.jit
+    def experiment(x):
+      def model(y, t):
+        return -x * y
+      history = odeint(model, 1., np.arange(0, 10, 0.1))
+      return history[-1]
+
+    gradfun = jax.value_and_grad(experiment)
+    t = np.arange(0., 1., 0.01)
+    h, g = jax.vmap(gradfun)(t)  # doesn't crash
+    ans = h[11], g[11]
+
+    expected_h = experiment(t[11])
+    expected_g = (experiment(t[11] + 1e-5) - expected_h) / 1e-5
+    expected = expected_h, expected_g
+
+    self.assertAllClose(ans, expected, check_dtypes=False, atol=1e-2, rtol=1e-2)
+
 
 if __name__ == '__main__':
   absltest.main()