Remove jax_custom_vjp_disable_shape_check config option

PiperOrigin-RevId: 845005784

Author: yashk2810

jax/_src/config.py

@@ -1819,13 +1819,6 @@ def _validate_default_device(val):
     upgrade=False,
     help='Temporary workaround to disable an error check in vmap-of-shmap.')
 
-# TODO(mattjj): remove once we land mutable array plumbing, or face great shame
-custom_vjp_disable_shape_check = 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.')
-
 mutable_array_checks = bool_state(
     name='jax_mutable_array_checks',
     default=True,

jax/_src/custom_derivatives.py

@@ -966,8 +966,7 @@ def append(x, d):
     else:
       if (not core.typecompat(a.to_tangent_aval(), a_ := core.get_aval(ct)) and
           not _ref_typecompat(a.to_tangent_aval(), a_) and
-          not (_temporary_dtype_exception(a, a_) or
-               _temporary_shape_exception(a, a_))):
+          not _temporary_dtype_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 "
@@ -990,9 +989,6 @@ 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.Primitive):
   multiple_results = True

tests/custom_api_test.py

@@ -2971,23 +2971,6 @@ 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)
-
-    with config.custom_vjp_disable_shape_check(True):
-      jax.grad(lambda x, y: foo(x, y).sum(), 1)(jnp.ones(3), jnp.ones(4))
-
   def test_bwd_rule_can_produce_list_or_tuple(self):
     @jax.custom_vjp
     def f(x, y):