Merge pull request #20599 from mattjj:temp-config-to-disable-custom-vjp-shape-check

PiperOrigin-RevId: 622224003

Author: jax authors

jax/_src/config.py

@@ -1390,6 +1390,13 @@ def _update_disable_jit_thread_local(val):
     upgrade=True,
     help='Enable eager-mode pmap when jax_disable_jit is activated.')
 
+# TODO(mattjj): remove once we land mutable array plumbing, or face great shame
+custom_vjp_disable_shape_check = define_bool_state(
+    name='jax_custom_vjp_disable_shape_check',
+    default=False,
+    upgrade=True,
+    help='Disable the check from #19009 to enable some custom_vjp hacks.')
+
 xla_runtime_errors = define_bool_state(
     name='jax_experimental_unsafe_xla_runtime_errors',
     default=False,

jax/_src/custom_derivatives.py

@@ -772,7 +772,8 @@ def append(x, d):
       results.append(Zero(ct.aval))
     else:
       if (not core.typecompat(a.at_least_vspace(), a_ := core.get_aval(ct))
-          and not _temporary_dtype_exception(a, a_)):
+          and not (_temporary_dtype_exception(a, a_) or
+                   _temporary_shape_exception(a, a_))):
         msg = ("Custom VJP bwd rule must produce an output with the same "
                "shape/dtypes as the args tuple of the primal function, but at "
                f"output{keystr(kp)} the bwd rule produced an output of "
@@ -790,6 +791,9 @@ def _temporary_dtype_exception(a, a_) -> bool:
              dtypes.issubdtype(a.dtype, dtypes.np.inexact)))
   return False
 
+# TODO(mattjj): remove both these exceptions to cotangent compatibility check
+def _temporary_shape_exception(a, a_) -> bool:
+  return config.custom_vjp_disable_shape_check.value
 
 class CustomVJPCallPrimitive(core.CallPrimitive):
   initial_style: core.Primitive

tests/api_test.py

@@ -9372,6 +9372,27 @@ def foo_bwd(_, g):
         r'output\[1\] the bwd rule produced an output of shape/dtype float..\[3\]'):
       jax.grad(lambda x, y: foo(x, y * y).sum(), 1)(jnp.ones(3), jnp.ones(4))
 
+  def test_bwd_rule_shape_mismatch_disable(self):
+    # TODO(mattjj): remove this test when the config option is removed
+    @jax.custom_vjp
+    def foo(x, y):
+      return x
+
+    def foo_fwd(x, y):
+      return x, None
+
+    def foo_bwd(_, g):
+      return jnp.zeros(3), jnp.zeros(3)
+
+    foo.defvjp(foo_fwd, foo_bwd)
+
+    try:
+      jax.config.update('jax_custom_vjp_disable_shape_check', True)
+      jax.grad(lambda x, y: foo(x, y).sum(), 1)(jnp.ones(3), jnp.ones(4))
+    finally:
+      jax.config.update('jax_custom_vjp_disable_shape_check', False)
+
+
 def transpose_unary(f, x_example):
   def transposed(y):
     x, = api.linear_transpose(f, x_example)(y)