Repository: google/jax

Issue:
host_callback doesn't work inside grad(odeint)
It reports "NotImplementedError: outfeed rewrite custom_vjp_call_jaxpr"

This would be quite useful to support because it would give us a way to get information about nature of the backward pass out of odeint, e.g., to facilitate debugging https://github.com/google/jax/issues/3993.

To reproduce:
```python
from jax.experimental.ode import odeint
from jax.experimental import host_callback
import jax.numpy as jnp
import jax

def f(x, t, k):
  x = host_callback.id_print(x)
  return -k * x

def loss(k=1.0):
  t = jnp.linspace(0, 0.001, num=2)
  xs = odeint(f, 1.0, t, k)
  return xs[-1]

loss(1.0)  # works
jax.grad(loss)(1.0)  # fails
```

The error message is:
```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-1-aea888293dfa> in <module>()
     14 
     15 loss(1.0)  # works
---> 16 jax.grad(loss)(1.0)  # fails

19 frames
/usr/local/lib/python3.6/dist-packages/jax/api.py in grad_f(*args, **kwargs)
    427   @wraps(fun, docstr=docstr, argnums=argnums)
    428   def grad_f(*args, **kwargs):
--> 429     _, g = value_and_grad_f(*args, **kwargs)
    430     return g
    431 

/usr/local/lib/python3.6/dist-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    491     dtype = dtypes.result_type(ans)
    492     tree_map(partial(_check_output_dtype_grad, holomorphic), ans)
--> 493     g = vjp_py(np.ones((), dtype=dtype))
    494     g = g[0] if isinstance(argnums, int) else g
    495     if not has_aux:

/usr/local/lib/python3.6/dist-packages/jax/api.py in _vjp_pullback_wrapper(cotangent_dtypes, io_tree, fun, py_args)
   1458              "match type of corresponding primal output ({})")
   1459       raise TypeError(msg.format(_dtype(a), dtype))
-> 1460   ans = fun(*args)
   1461   return tree_unflatten(out_tree, ans)
   1462 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    115     cts = tuple(map(ignore_consts, cts, pvals))
    116     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--> 117     arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
    118     return map(instantiate_zeros, arg_cts)
    119 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in backward_pass(jaxpr, consts, primals_in, cotangents_in)
    201         call_jaxpr, params = core.extract_call_jaxpr(eqn.primitive, eqn.params)
    202         cts_out = get_primitive_transpose(eqn.primitive)(
--> 203             params, call_jaxpr, invals, cts_in, cts_in_avals)
    204       else:
    205         cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,

/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py in call_transpose(primitive, params, call_jaxpr, args, ct, _)
    486     new_params = update_params(new_params, map(is_undefined_primal, args),
    487                                [type(x) is not Zero for x in ct])
--> 488   out_flat = primitive.bind(fun, *all_args, **new_params)
    489   return tree_unflatten(out_tree(), out_flat)
    490 primitive_transposes[core.call_p] = partial(call_transpose, call_p)

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, fun, *args, **params)
   1132 
   1133   def bind(self, fun, *args, **params):
-> 1134     return call_bind(self, fun, *args, **params)
   1135 
   1136   def process(self, trace, fun, tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1121   if top_trace is None:
   1122     with new_sublevel():
-> 1123       outs = primitive.impl(fun, *args, **params)
   1124   else:
   1125     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    525 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
    526   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
--> 527                                *unsafe_map(arg_spec, args))
    528   try:
    529     return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    222       fun.populate_stores(stores)
    223     else:
--> 224       ans = call(fun, *args)
    225       cache[key] = (ans, fun.stores)
    226     return ans

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    598         fun, pvals, instantiate=False, stage_out=True, bottom=True)
    599   map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))
--> 600   jaxpr = apply_outfeed_rewriter(jaxpr)
    601 
    602   nreps = jaxpr_replicas(jaxpr)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in apply_outfeed_rewriter(jaxpr)
    181 def apply_outfeed_rewriter(jaxpr: core.Jaxpr) -> core.Jaxpr:
    182   if outfeed_rewriter is not None:
--> 183     return outfeed_rewriter(jaxpr)
    184   else:
    185     return jaxpr

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in <lambda>(j)
    752 
    753 
--> 754 xla.outfeed_rewriter = lambda j: _rewrite_jaxpr(j, False, False)
    755 
    756 

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
    553     else:
    554       output_token_var = mk_new_var(core.abstract_token)
--> 555       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var, mk_new_var)
    556       last_token_var = output_token_var
    557 

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, mk_new_var)
    617         input_token_var
    618     ] + eqn.invars[nr_const_and_carry:]
--> 619     new_jaxpr = _rewrite_typed_jaxpr(carry_jaxpr, True, True)
    620     # The rewrite has put the token at end, it has to be at end of carry
    621     new_jaxpr_invars = new_jaxpr.jaxpr.invars

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_typed_jaxpr(tjaxpr, has_input_token, has_output_token)
    524     has_output_token: bool) -> core.TypedJaxpr:
    525   """Rewrites a TypedJaxpr to thread the token, if needed."""
--> 526   new_jaxpr = _rewrite_jaxpr(tjaxpr.jaxpr, has_input_token, has_output_token)
    527   return _mk_typed_jaxpr(new_jaxpr, tjaxpr.literals)
    528 

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
    553     else:
    554       output_token_var = mk_new_var(core.abstract_token)
--> 555       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var, mk_new_var)
    556       last_token_var = output_token_var
    557 

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, mk_new_var)
    654                 eqn.params,
    655                 call_jaxpr=_rewrite_jaxpr(call_jaxpr, True,
--> 656                                           True)), eqn.source_info))
    657   else:
    658     raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
    553     else:
    554       output_token_var = mk_new_var(core.abstract_token)
--> 555       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var, mk_new_var)
    556       last_token_var = output_token_var
    557 

/usr/local/lib/python3.6/dist-packages/jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, mk_new_var)
    656                                           True)), eqn.source_info))
    657   else:
--> 658     raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")
    659 
    660 

NotImplementedError: outfeed rewrite custom_vjp_call_jaxpr
```


Discussion:
@gnecula want to take this one?

Resolving patch:
diff --git a/jax/experimental/host_callback.py b/jax/experimental/host_callback.py
--- a/jax/experimental/host_callback.py
+++ b/jax/experimental/host_callback.py
@@ -136,6 +136,7 @@ def power3(x):
 
 from jax import api
 from jax import core
+from jax import custom_derivatives
 from jax import lax
 from jax.lib import pytree
 from jax.interpreters import ad, xla, batching, masking
@@ -654,6 +655,39 @@ def _rewrite_eqn(eqn: core.JaxprEqn, eqns: List[core.JaxprEqn],
                 eqn.params,
                 call_jaxpr=_rewrite_jaxpr(call_jaxpr, True,
                                           True)), eqn.source_info))
+  elif eqn.primitive is custom_derivatives.custom_jvp_call_jaxpr_p:
+    fun_jaxpr = eqn.params["fun_jaxpr"]
+    new_invars = [*eqn.invars, input_token_var]
+    def unreachable_thunk():
+      assert False, "Should not be reached"
+    eqns.append(
+        core.new_jaxpr_eqn(
+            new_invars, eqn.outvars + [output_token_var], eqn.primitive,
+            dict(
+                eqn.params,
+                fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
+                jvp_jaxpr_thunk=unreachable_thunk
+            ),
+            eqn.source_info))
+  elif eqn.primitive is custom_derivatives.custom_vjp_call_jaxpr_p:
+    fun_jaxpr = eqn.params["fun_jaxpr"]
+    new_invars = [*eqn.invars, input_token_var]
+    def unreachable_thunk():
+      assert False, "Should not be reached"
+    eqns.append(
+        core.new_jaxpr_eqn(
+            new_invars, eqn.outvars + [output_token_var], eqn.primitive,
+            dict(
+                eqn.params,
+                fun_jaxpr=_rewrite_typed_jaxpr(fun_jaxpr, True, True),
+                fwd_jaxpr_thunk=unreachable_thunk,
+                # The following are illegal values for the parameters, they
+                # should not be needed because this rewrite is just before
+                # compilation to XLA, which does not use those parameters.
+                bwd="illegal param",
+                out_trees="illegal param"
+            ),
+            eqn.source_info))
   else:
     raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")
 


Test patch:
diff --git a/tests/host_callback_test.py b/tests/host_callback_test.py
--- a/tests/host_callback_test.py
+++ b/tests/host_callback_test.py
@@ -90,6 +90,10 @@ def repl_floats(match_group):
   what = re.sub(r"\-?\d*\.[\-\def]*", repl_floats, what)
   what = re.sub(r"output_stream=[^\]\n]*", "", what)
   what = re.sub(r"threshold=[^\]\n]*", "", what)
+  what = re.sub(r"bwd=[^\]\n]*", "", what)
+  what = re.sub(r"out_trees=[^\]\n]*", "", what)
+  what = re.sub(r"fwd_jaxpr_thunk=[^\]\n]*", "", what)
+  what = re.sub(r"jvp_jaxpr_thunk=[^\]\n]*", "", what)
   # Empty lines
   what = re.sub(r"^\s*\n", "", what, flags=re.MULTILINE)
   def repl_func(match_group):
@@ -916,6 +920,94 @@ def test_pmap(self):
     expected_res = jnp.stack([fun1_equiv(2. + a) for a in range(api.local_device_count())])
     self.assertAllClose(expected_res, res, check_dtypes=False)
 
+  def test_scan_custom_jvp(self):
+    """custom JVP, inside scan.
+    This exercises the custom_jvp_call_jaxpr primitives."""
+    @api.custom_jvp
+    def f(x):
+      return x * hcb.id_print(x, output_stream=testing_stream, what="x")
+
+    @f.defjvp
+    def f_jvp(primals, tangents):
+      x, = primals
+      x_dot, = tangents
+      primal_out = f(x)
+      tangent_out = 3. * x * hcb.id_print(x_dot, output_stream=testing_stream, what="x_dot")
+      return primal_out, tangent_out
+
+    def g(x):
+      # Sum f(x_i)
+      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
+                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
+                      x)[0]
+
+    arg = np.full((2,), 0.7)
+    self.assertAllClose(0.7 * 0.7 * 2, g(arg))
+    hcb.barrier_wait()
+    self.assertMultiLineStrippedEqual("""
+    what: x
+    0.7
+    what: x
+    0.7""", testing_stream.output)
+    testing_stream.reset()
+
+    self.assertAllClose(np.array([2.1, 2.1]), api.grad(g)(arg), check_dtypes=False)
+    hcb.barrier_wait()
+    self.assertMultiLineStrippedEqual("""
+    what: x
+    0.7
+    what: x
+    0.7
+    transforms: ({'name': 'transpose'},) what: x_dot
+    2.1
+    transforms: ({'name': 'transpose'},) what: x_dot
+    2.1""", testing_stream.output)
+
+  def test_scan_custom_vjp(self):
+    """custom VJP, inside scan.
+    This exercises the custom_vjp_call_jaxpr primitives."""
+    @api.custom_vjp
+    def f(x):
+      return x * hcb.id_print(x, output_stream=testing_stream, what="x")
+
+    # f_fwd: a -> (b, residual)
+    def f_fwd(x):
+      return f(x), 3. * x
+    # f_bwd: (residual, CT b) -> [CT a]
+    def f_bwd(residual, ct_b):
+      return residual * hcb.id_print(ct_b, output_stream=testing_stream, what="ct_b"),
+
+    f.defvjp(f_fwd, f_bwd)
+
+    def g(x):
+      # Sum f(x_i)
+      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
+                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
+                      x)[0]
+
+    arg = np.full((2,), 0.7)
+
+    self.assertAllClose(0.7 * 0.7 * 2, g(arg))
+    hcb.barrier_wait()
+    self.assertMultiLineStrippedEqual("""
+      what: x
+      0.7
+      what: x
+      0.7""", testing_stream.output)
+    testing_stream.reset()
+
+    self.assertAllClose(np.array([2.1, 2.1]), api.grad(g)(arg), check_dtypes=False)
+    hcb.barrier_wait()
+    self.assertMultiLineStrippedEqual("""
+      what: x
+      0.7
+      what: x
+      0.7
+      what: ct_b
+      1.
+      what: ct_b
+      1.""", testing_stream.output)
+
   def test_mask(self):
     # TODO(necula)
     raise SkipTest("masking has regressed")
@@ -1186,9 +1278,176 @@ def func(x):
               linear=(False, False, False, False, False)
               num_carry=3
               num_consts=1
-              reverse=False ] b 1 2 f a
+              reverse=False
+              unroll=1 ] b 1 2 f a
   in (c, d, e, g) }""", func, [y])
 
+  def test_scan_custom_jvp(self):
+    """custom JVP, inside scan.
+    This exercises the custom_jvp_call_jaxpr primitives."""
+    @api.custom_jvp
+    def f(x):
+      return x * hcb.id_print(x)
+
+    @f.defjvp
+    def f_jvp(primals, tangents):
+      x, = primals
+      x_dot, = tangents
+      primal_out = f(x)
+      tangent_out = 3. * x * hcb.id_print(x_dot)
+      return primal_out, tangent_out
+
+    def g(x):
+      # Sum f(x_i)
+      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
+                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
+                      x)[0]
+
+    arg = np.full((5,), 0.7)
+    self.assertRewrite("""
+    { lambda  ; a c.
+      let b d _ = scan[ jaxpr={ lambda  ; a e b.
+                                let c f = custom_jvp_call_jaxpr[ fun_jaxpr={ lambda  ; a d.
+                                                                             let b e = id_tap[ arg_treedef_=*
+                                                                                               has_token_=True
+                                                                                               nr_tapped_args_=1
+                                                                                               tap_func_=_print
+                                                                                               ] a d
+                                                                                 c = mul a b
+                                                                             in (c, e) }
+                                                                 ] b e
+                                    d = add a c
+                                in (d, f, 0.00) }
+                        length=5
+                        linear=(False, False, False)
+                        num_carry=2
+                        num_consts=0
+                        reverse=False
+                        unroll=1 ] 0.00 c a
+      in (b, d) }""", g, [arg])
+    self.assertRewrite("""
+    { lambda  ; a d.
+      let _ _ e _ b =
+            scan[ jaxpr={ lambda  ; a b h c d.
+                          let e i = custom_jvp_call_jaxpr[ fun_jaxpr={ lambda  ; a d.
+                                                                       let b e = id_tap[ arg_treedef_=*
+                                                                                         has_token_=True
+                                                                                         nr_tapped_args_=1
+                                                                                         tap_func_=_print
+                                                                                         ] a d
+                                                                           c = mul a b
+                                                                       in (c, e) }
+                                                           ] c h
+                              f = add a e
+                              g = mul c 3.00
+                          in (f, *, i, 0.00, g) }
+                  length=5
+                  linear=(False, True, False, True, False)
+                  num_carry=3
+                  num_consts=0
+                  reverse=False
+                  unroll=1 ] 0.00 * d a *
+          _ _ f _ c =
+            scan[ jaxpr={ lambda  ; a b g c d.
+                          let e = mul b d
+                              f h = id_tap[ arg_treedef_=*
+                                            has_token_=True
+                                            nr_tapped_args_=1
+                                            tap_func_=_print
+                                            transforms=(('transpose',),) ] e g
+                          in (*, b, h, *, f) }
+                  length=5
+                  linear=(True, True, True, False, False)
+                  num_carry=3
+                  num_consts=0
+                  reverse=True
+                  unroll=1 ] * 1.00 e * b
+      in (c, f) }""", api.grad(g), [arg])
+
+  def test_scan_custom_vjp(self):
+    """custom VJP, inside scan.
+    This exercises the custom_vjp_call_jaxpr primitives."""
+    @api.custom_vjp
+    def f(x):
+      return x * hcb.id_print(x)
+
+    # f_fwd: a -> (b, residual)
+    def f_fwd(x):
+      return f(x), 3. * x
+    # f_bwd: (residual, CT b) -> [CT a]
+    def f_bwd(residual, ct_b):
+      return residual * hcb.id_print(ct_b),
+
+    f.defvjp(f_fwd, f_bwd)
+
+    def g(x):
+      # Sum f(x_i)
+      return lax.scan(lambda carry, inp: (carry + f(inp), 0.),
+                      np.full(x.shape[1:], 0.),  # Like x w/o leading dim
+                      x)[0]
+
+    arg = np.full((2,), 0.7)
+    self.assertRewrite("""
+    { lambda  ; a c.
+      let b d _ = scan[ jaxpr={ lambda  ; a e b.
+                                let c f = custom_vjp_call_jaxpr[
+                                                                 fun_jaxpr={ lambda  ; a d.
+                                                                             let b e = id_tap[ arg_treedef_=*
+                                                                                               has_token_=True
+                                                                                               nr_tapped_args_=1
+                                                                                               tap_func_=_print
+                                                                                               ] a d
+                                                                                 c = mul a b
+                                                                             in (c, e) }
+                                                                 ] b e
+                                    d = add a c
+                                in (d, f, 0.00) }
+                        length=2
+                        linear=(False, False, False)
+                        num_carry=2
+                        num_consts=0
+                        reverse=False
+                        unroll=1 ] 0.00 c a
+      in (b, d) }""", g, [arg])
+    self.assertRewrite("""
+    { lambda  ; a d.
+      let _ _ e _ b =
+            scan[ jaxpr={ lambda  ; a b h c d.
+                          let e i = custom_vjp_call_jaxpr[
+                                                           fun_jaxpr={ lambda  ; a d.
+                                                                       let b e = id_tap[ arg_treedef_=*
+                                                                                         has_token_=True
+                                                                                         nr_tapped_args_=1
+                                                                                         tap_func_=_print
+                                                                                         ] a d
+                                                                           c = mul a b
+                                                                       in (c, e) }
+                                                           ] c h
+                              f = add a e
+                              g = mul c 3.00
+                          in (f, *, i, 0.00, g) }
+                  length=2
+                  linear=(False, True, False, True, False)
+                  num_carry=3
+                  num_consts=0
+                  reverse=False
+                  unroll=1 ] 0.00 * d a *
+          _ _ f _ c =
+            scan[ jaxpr={ lambda  ; a b g c d.
+                          let e h = id_tap[ arg_treedef_=*
+                                            has_token_=True
+                                            nr_tapped_args_=1
+                                            tap_func_=_print
+                                            ] b g
+                              f = mul d e
+                          in (*, b, h, *, f) }
+                  length=2
+                  linear=(True, True, True, False, False)
+                  num_carry=3
+                  num_consts=0
+                  reverse=True
+                  unroll=1 ] * 1.00 e * b
+      in (c, f) }""", api.grad(g), [arg])
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())