Repository: google/jax
Pull request #23000 (2024-08-12 12:36:40)

Fix bug in custom_vjp with optimize_remat and custom_vmap
When used with a `custom_vmap` that introduces a new const the previous implementation of `optimize_remat` would error in its DCE rule because of unexpected consts when closing the fwd jaxpr. This shouldn't have ever been hit, but there was a bug in the batching rule for `remat_opt_p` where we weren't properly converting constvars to invars. This fixes this bug and should unbreak internal users.


Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 6dbf7dc11..9a6253b1b 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -1523,6 +1523,9 @@ def _remat_opt_vmap(
   batched_fwd_jaxpr, out_batched = batching.batch_jaxpr(
       fwd_jaxpr, axis_size, in_batched, False,
       axis_name, spmd_axis_name, main_type)
+  extra_consts = batched_fwd_jaxpr.consts
+  batched_fwd_jaxpr = pe.close_jaxpr(
+      pe.convert_constvars_jaxpr(batched_fwd_jaxpr.jaxpr))
   out_dims = [0 if b else not_mapped for b in out_batched]
 
   _, prim_batched = split_list(in_batched, [num_consts])
@@ -1535,7 +1538,8 @@ def _remat_opt_vmap(
         main_type)
     return batched_fun_jaxpr.jaxpr, batched_fun_jaxpr.consts
 
-  batched_outs = remat_opt_p.bind(*args, num_consts=num_consts,
+  batched_outs = remat_opt_p.bind(*extra_consts, *args,
+                                  num_consts=num_consts + len(extra_consts),
                                   num_res=num_res,
                                   fwd_jaxpr=batched_fwd_jaxpr,
                                   fun_jaxpr_thunk=batched_fun_jaxpr_thunk)
@@ -1603,6 +1607,7 @@ def _remat_opt_dce(used_outs: list[bool], eqn: core.JaxprEqn):
     instantiate += [True] * (len(eqn.invars) - eqn.params["num_consts"])
     new_jaxpr, used_ins = pe.dce_jaxpr(eqn.params["fwd_jaxpr"].jaxpr, used_outs,
                                        instantiate=instantiate)
+    assert not new_jaxpr.constvars
     closed_jaxpr = pe.close_jaxpr(new_jaxpr)
     invars = [v for used, v in zip(used_ins, eqn.invars) if used]
     new_params = dict(eqn.params)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index cba74bee5..eee492443 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -9779,6 +9779,31 @@ class CustomVJPTest(jtu.JaxTestCase):
     x, y = 3.2, 1.0
     jax.grad(f)(x, y)  # Doesn't error
 
+  def test_optimize_remat_custom_vmap(self):
+    # See https://github.com/google/jax/pull/23000
+    @jax.custom_vjp
+    def f(x, y):
+      return jnp.sin(x) * y
+
+    @jax.custom_batching.custom_vmap
+    def f_fwd(x, y):
+      return f(x, y), (jnp.cos(x), jnp.sin(x), y)
+
+    @f_fwd.def_vmap
+    def f_fwd_vmap(_, in_batched, x, y):
+      # Insert a new const here to test the optimize_remat batching rule.
+      out = np.array([2.0])*f(x, y)
+      out_batched = (True, (True, True, True))
+      return (out, (jnp.cos(x), jnp.sin(x), y)), out_batched
+
+    def f_bwd(res, g):
+      cos_x, sin_x, y = res
+      return (cos_x * g * y, sin_x * g)
+
+    f.defvjp(f_fwd, f_bwd, optimize_remat=True)
+    x, y = jnp.linspace(0.0, 1.0, 5), jnp.linspace(2.0, 5.0, 5)
+    jax.jit(jax.vmap(jax.grad(f)))(x, y)  # Doesn't error
+
 
 def transpose_unary(f, x_example):
   def transposed(y):