Repository: google/jax
Issue: Host callbacks + remat causes NotImplementedError

Would it be possible to add remat to host_callback's? I am getting the following error: `NotImplementedError: outfeed rewrite remat_call`.

Repro:
```python
from jax.experimental import host_callback
import jax.numpy as jnp
import jax
from jax import lax

def f(k, x):
  x = host_callback.id_print(k+x)
  return -k * x

def loss(k=1.0):
  return lax.fori_loop(0, 1, jax.remat(f), 0)(k)
  
loss(1)
```

```
Stack:
<ipython-input-72-12c7a8d25b00> in <module>()
     11   return lax.fori_loop(0, 1, jax.remat(f), 0)(k)
     12 
---> 13 loss(1)

15 frames
<ipython-input-72-12c7a8d25b00> in loss(k)
      9 
     10 def loss(k=1.0):
---> 11   return lax.fori_loop(0, 1, jax.remat(f), 0)(k)
     12 
     13 loss(1)

jax/lax/lax_control_flow.py in fori_loop(lower, upper, body_fun, init_val)
    204   else:
    205     _, _, result = while_loop(_fori_cond_fun, _fori_body_fun(body_fun),
--> 206                               (lower, upper, init_val))
    207   return result
    208 

jax/lax/lax_control_flow.py in while_loop(cond_fun, body_fun, init_val)
    297   outs = while_p.bind(*itertools.chain(cond_consts, body_consts, init_vals),
    298                       cond_nconsts=len(cond_consts), cond_jaxpr=cond_jaxpr,
--> 299                       body_nconsts=len(body_consts), body_jaxpr=body_jaxpr)
    300   return tree_unflatten(body_tree, outs)
    301 

jax/core.py in bind(self, *args, **params)
    261     top_trace = find_top_trace(args)
    262     tracers = map(top_trace.full_raise, args)
--> 263     out = top_trace.process_primitive(self, tracers, params)
    264     return map(full_lower, out) if self.multiple_results else full_lower(out)
    265 

jax/core.py in process_primitive(self, primitive, tracers, params)
    570 
    571   def process_primitive(self, primitive, tracers, params):
--> 572     return primitive.impl(*tracers, **params)
    573 
    574   def process_call(self, primitive, f, tracers, params):

jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    231 def apply_primitive(prim, *args, **params):
    232   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 233   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
    234   return compiled_fun(*args)
    235 

jax/interpreters/xla.py in xla_primitive_callable(prim, *arg_specs, **params)
    255       return prim.bind(*args, **params)
    256     return _xla_callable(lu.wrap_init(prim_fun), device, None, "prim", donated_invars,
--> 257                          *arg_specs)
    258   aval_out = prim.abstract_eval(*avals, **params)
    259   if not prim.multiple_results:

jax/linear_util.py in memoized_fun(fun, *args)
    245       fun.populate_stores(stores)
    246     else:
--> 247       ans = call(fun, *args)
    248       cache[key] = (ans, fun.stores)
    249 

jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    638         fun, pvals, instantiate=False, stage_out=True, bottom=True)  # type: ignore
    639   map(prefetch, it.chain(consts, jaxpr_literals(jaxpr)))
--> 640   jaxpr = apply_outfeed_rewriter(jaxpr)
    641 
    642   nreps = jaxpr_replicas(jaxpr)

jax/interpreters/xla.py in apply_outfeed_rewriter(jaxpr)
    190 def apply_outfeed_rewriter(jaxpr: core.Jaxpr) -> core.Jaxpr:
    191   if outfeed_rewriter is not None:
--> 192     return outfeed_rewriter(jaxpr)
    193   else:
    194     return jaxpr

jax/experimental/host_callback.py in <lambda>(j)
    831 
    832 
--> 833 xla.outfeed_rewriter = lambda j: _rewrite_jaxpr(j, False, False)
    834 
    835 

jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
    601     else:
    602       output_token_var = mk_new_var(core.abstract_token)
--> 603       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var, mk_new_var)
    604       last_token_var = output_token_var
    605 

jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, mk_new_var)
    639             dict(
    640                 eqn.params,
--> 641                 body_jaxpr=_rewrite_closed_jaxpr(body_jaxpr, True, True),
    642                 cond_jaxpr=_rewrite_closed_jaxpr(cond_jaxpr, True,
    643                                                 False)), eqn.source_info))
jax/experimental/host_callback.py in _rewrite_closed_jaxpr(cjaxpr, has_input_token, has_output_token)
    571     has_output_token: bool) -> core.ClosedJaxpr:
    572   """Rewrites a ClosedJaxpr to thread the token, if needed."""
--> 573   new_jaxpr = _rewrite_jaxpr(cjaxpr.jaxpr, has_input_token, has_output_token)
    574   return core.ClosedJaxpr(new_jaxpr, cjaxpr.consts)
    575 

py/jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
    601     else:
    602       output_token_var = mk_new_var(core.abstract_token)
--> 603       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var, mk_new_var)
    604       last_token_var = output_token_var
    605 

jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, mk_new_var)
    738             eqn.source_info))
    739   else:
--> 740     raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")
    741 
    742 

NotImplementedError: outfeed rewrite remat_call
```

Changed files:
jax/interpreters/xla.py, jax/experimental/host_callback.py, tests/host_callback_test.py

diff --git a/jax/experimental/host_callback.py b/jax/experimental/host_callback.py
index fd1b3ce73..26a962a68 100644
--- a/jax/experimental/host_callback.py
+++ b/jax/experimental/host_callback.py
@@ -613,6 +613,7 @@ def _rewrite_eqn(eqn: core.JaxprEqn, eqns: List[core.JaxprEqn],
                  mk_new_var: Callable[[core.AbstractValue], core.Var]):
   """Rewrite an `eqn` and append equations to `eqns`.
 
+  This is only called if the current primitive uses outfeed.
   Assume that the current token is in `input_token_var` and the resulting
   token must end in `output_token_var`.
   """
@@ -698,11 +699,21 @@ def _rewrite_eqn(eqn: core.JaxprEqn, eqns: List[core.JaxprEqn],
             eqn.primitive,
             dict(
                 eqn.params,
-                call_jaxpr=_rewrite_jaxpr(call_jaxpr, True,
-                                          True),
+                call_jaxpr=_rewrite_jaxpr(call_jaxpr, True, True),
                 donated_invars=eqn.params["donated_invars"] + (False,)
             ),
           eqn.source_info))
+  elif eqn.primitive is pe.remat_call_p:
+    call_jaxpr = cast(core.Jaxpr, eqn.params["call_jaxpr"])
+    eqns.append(
+        core.new_jaxpr_eqn(
+            eqn.invars + [input_token_var], eqn.outvars + [output_token_var],
+            eqn.primitive,
+            dict(
+                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]
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 154520ad2..9412f7d16 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -1305,9 +1305,13 @@ def _remat_translation_rule(c, axis_env, in_nodes,
   xb.parameter(dummy_subc, 0, c.get_shape(false_op), replicated=[])
 
   def zeros(xla_shape):
-    shape, dtype = xla_shape.dimensions(), xla_shape.numpy_dtype()
-    zero = xb.constant(dummy_subc, np.array(0, dtype=dtype))
-    return xops.Broadcast(zero, shape)
+    if xla_shape.is_array():
+      shape, dtype = xla_shape.dimensions(), xla_shape.numpy_dtype()
+      zero = xb.constant(dummy_subc, np.array(0, dtype=dtype))
+      return xops.Broadcast(zero, shape)
+    else:
+      # It is a token
+      return xops.CreateToken(dummy_subc)
   out_nodes = [zeros(s) for s in out_node_shapes]
   dummy_subc = dummy_subc.build(xops.Tuple(dummy_subc, out_nodes))
 
diff --git a/tests/host_callback_test.py b/tests/host_callback_test.py
index 055905d44..d6fb98550 100644
--- a/tests/host_callback_test.py
+++ b/tests/host_callback_test.py
@@ -149,7 +149,6 @@ class HostCallbackTest(jtu.JaxTestCase):
       x1, y1 = hcb.id_print((x * 2., x * 3.), output_stream=testing_stream)
       return x1 + y1
 
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(func2)(3.)))
     self.assertEqual(3. * (2. + 3.), func2(3.))
     hcb.barrier_wait()
 
@@ -220,8 +219,6 @@ class HostCallbackTest(jtu.JaxTestCase):
     def func():  # jitted function does not take arguments
       return hcb.id_print(42, output_stream=testing_stream)
 
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(api.jit(func))()))
-
     self.assertAllClose(42, api.jit(func)())
     hcb.barrier_wait()
     assertMultiLineStrippedEqual(self, """
@@ -232,8 +229,6 @@ class HostCallbackTest(jtu.JaxTestCase):
     def func(x1, x2):
       return hcb.id_print(x1 + x2, output_stream=testing_stream)
 
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(api.jit(func))(40, 2)))
-
     self.assertAllClose(42, api.jit(func)(40, 2))
     hcb.barrier_wait()
     assertMultiLineStrippedEqual(self, """
@@ -244,8 +239,6 @@ class HostCallbackTest(jtu.JaxTestCase):
     def func(x):
       return hcb.id_print(42, result=x, output_stream=testing_stream)
 
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(api.jit(func))(5)))
-
     self.assertAllClose(5, api.jit(func)(5))
     hcb.barrier_wait()
     assertMultiLineStrippedEqual(self, """
@@ -727,7 +720,6 @@ class HostCallbackTest(jtu.JaxTestCase):
 
   def test_jvp(self):
     jvp_fun1 = lambda x, xt: api.jvp(fun1, (x,), (xt,))
-    #assertMultiLineStrippedEqual(self, "")
     res_primals, res_tangents = jvp_fun1(jnp.float32(5.), jnp.float32(0.1))
     self.assertAllClose(100., res_primals, check_dtypes=False)
     self.assertAllClose(4., res_tangents, check_dtypes=False)
@@ -790,7 +782,6 @@ class HostCallbackTest(jtu.JaxTestCase):
       return x * hcb.id_print(y * 3., what="y * 3",
                               output_stream=testing_stream)
     grad_func = api.grad(func)
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(grad_func)(5.)))
 
     res_grad = grad_func(jnp.float32(5.))
     self.assertAllClose(2. * 5. * 6., res_grad, check_dtypes=False)
@@ -837,7 +828,6 @@ class HostCallbackTest(jtu.JaxTestCase):
   def test_vmap(self):
     vmap_fun1 = api.vmap(fun1)
     vargs = jnp.array([jnp.float32(4.), jnp.float32(5.)])
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(vmap_fun1)(vargs)))
     vmap_fun1(vargs)
     hcb.barrier_wait()
     assertMultiLineStrippedEqual(self, """
@@ -856,7 +846,6 @@ class HostCallbackTest(jtu.JaxTestCase):
 
     vmap_func = api.vmap(func)
     vargs = jnp.array([jnp.float32(4.), jnp.float32(5.)])
-    #assertMultiLineStrippedEqual(self, "", str(api.make_jaxpr(vmap_func)(vargs)))
     _ = vmap_func(vargs)
     hcb.barrier_wait()
     assertMultiLineStrippedEqual(self, """
@@ -1177,6 +1166,20 @@ class HostCallbackTest(jtu.JaxTestCase):
 
     api.grad(loss)(1.0)  # should not fail
 
+  def test_remat(self):
+    def f(i, k):
+      x = hcb.id_print(k + i, output_stream=testing_stream)
+      return k * x
+
+    def loss(k):
+      return lax.fori_loop(0, 2, api.remat(f), k)
+    print(loss(3))
+    hcb.barrier_wait()
+    expected = """
+      3
+      10"""
+    self.assertMultiLineStrippedEqual(expected, testing_stream.output)
+
 
 class OutfeedRewriterTest(jtu.JaxTestCase):
 
@@ -1184,9 +1187,12 @@ class OutfeedRewriterTest(jtu.JaxTestCase):
                     has_input_token=True, has_output_token=True):
     """Check that the rewrite of func(*args) matches expected."""
     jaxpr = api.make_jaxpr(func)(*args)
-    # TODO: re-enable when we change the host_callback rewriter
-    #rewritten = hcb._rewrite_closed_jaxpr(jaxpr,
-    #                                      has_input_token, has_output_token)
+    rewritten = hcb._rewrite_closed_jaxpr(jaxpr,  # noqa: F841
+                                          has_input_token, has_output_token)
+    # Since it is somewhat annoying to update the Jaxpr assertions when we change
+    # the Jaxpr printing, we do not check these by default. It is recommended that
+    # before making changes to the code generation and Jaxpr rewriting, turn on
+    # the checking, update the expected Jaxpr, and then make the changes.
     #assertMultiLineStrippedEqual(self, expected, str(rewritten))
     del jaxpr
 
@@ -1543,5 +1549,38 @@ class OutfeedRewriterTest(jtu.JaxTestCase):
                       unroll=1 ] * 1.00 e * b
           in (c, f) }""", api.grad(g), [arg])
 
+  def test_remat_loop(self):
+    def f(k, x):
+      x = hcb.id_print(k + x)
+      return -k * x
+
+    def loss(k):
+      return lax.fori_loop(0, 1, api.remat(f), k)
+
+    self.assertRewrite("""
+        { lambda  ; a c.
+          let _ _ b d =
+                while[ body_jaxpr={ lambda  ; a b c f.
+                                    let d = add a 1
+                                        e g = remat_call[ call_jaxpr={ lambda  ; a b g.
+                                                                       let c = add a b
+                                                                           d h = id_tap[ arg_treedef_=*
+                                                                                         has_token_=True
+                                                                                         nr_tapped_args_=1
+                                                                                         tap_func_=_print  ] c g
+                                                                           e = neg a
+                                                                           f = mul e d
+                                                                       in (f, h) }
+                                                          concrete=False
+                                                          name=f ] a c f
+                                    in (d, b, e, g) }
+                       body_nconsts=0
+                       cond_jaxpr={ lambda  ; a b c e.
+                                    let d = lt a b
+                                    in (d,) }
+                       cond_nconsts=0 ] 0 1 a c
+          in (b, d) }""", loss, [2])
+
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())