Using remat with host_callback is broken

This throws an error:

import jax
from jax.experimental import host_callback

def f1(x):
  x = host_callback.id_tap(print, x)
  return x

def f2(x, _):
  f1_r = jax.remat(f1)
  x = f1_r(x)
  return x, 1

def forward(x):
  x = jax.lax.scan(f2, x, None, length=2)
  return x[0]

jax.jit(jax.grad(forward))(1.2)

...
# /usr/local/lib/python3.7/dist-packages/jax/experimental/host_callback.py in _outside_call_partial_eval_rule(trace, *args, **params)
    985   primals, tangents = util.split_list(args, [nr_primals])
    986   c_primals_tapped, _ = util.split_list(consts, [nr_primals])
--> 987   assert all([c is not None for c in c_primals_tapped])
    988 
    989   prims, _ = params["arg_treedef"].unflatten(args)

AssertionError: 

It runs fine if the line x = host_callback.id_tap(print, x) is removed, or if no lax.scan is used.

[sbodenstein]

Actually, the scan doesn’t seem necessary to cause the failure, this also works:

@jax.remat
def f(x):
  host_callback.id_tap(print, x)
  return x

jax.grad(f)(2.)

[gnecula]

Sadly, I have progress to report, except that I have reproduced the issue and saw that the error arises in the partial eval rule for the outside_call primitive. I wrote that code but with low confidence, and this issue is proof that there is still something I do not understand.

I can try to dig more into it, but it may be more useful if @mattjj finds some time to take a look also.