Remove config.vjp3 and vjp3 API since it's now replaced by jax.vjp

PiperOrigin-RevId: 843971635

Author: yashk2810

jax/_src/api.py

@@ -2212,31 +2212,6 @@ def vjp(
       fun, debug_info=debug_info("vjp", fun, primals, {}))
   return _vjp(wrapped_fun, *primals, has_aux=has_aux)
 
-def _vjp(fun: lu.WrappedFun, *primals, has_aux=False):
-  """Variant of vjp() that takes an lu.WrappedFun."""
-  if config.vjp3.value:
-    return _vjp3(fun, *primals, has_aux=has_aux)
-  primals_flat, in_tree = tree_flatten(primals)
-  primals_flat = [canonicalize_value(v) if not isinstance(v, core.Tracer) else v
-                  for v in primals_flat]
-  for arg in primals_flat: dispatch.check_arg(arg)
-  if not has_aux:
-    flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-    out_primals, vjp = ad.vjp(flat_fun, primals_flat)
-    out_tree = out_tree()
-  else:
-    flat_fun, out_aux_trees = flatten_fun_nokwargs2(fun, in_tree)
-    out_primals, vjp, aux = ad.vjp(flat_fun, primals_flat, has_aux=True)
-    out_tree, aux_tree = out_aux_trees()
-  out_primal_avals = map(shaped_abstractify, out_primals)
-  out_primal_py = tree_unflatten(out_tree, out_primals)
-  vjp_py = Partial(partial(_vjp_pullback_wrapper, fun.__name__,
-                           out_primal_avals, (out_tree, in_tree)), vjp)
-  if not has_aux:
-    return out_primal_py, vjp_py
-  else:
-    return out_primal_py, vjp_py, tree_unflatten(aux_tree, aux)
-
 @partial(api_boundary, repro_api_name="jax.experimental.saved_input_vjp")
 def saved_input_vjp(f: Callable, which: Sequence[bool], *primals,
                     allow_unused: bool = True, allow_opaque: bool = True):
@@ -2332,12 +2307,7 @@ class RSpec:
 si_vjp = saved_input_vjp
 
 
-def vjp3(f, *primals, has_aux=False):
-  dbg = debug_info("vjp", f, primals, {})
-  fun = lu.wrap_init(f, debug_info=dbg)
-  return _vjp3(fun, *primals, has_aux=has_aux)
-
-def _vjp3(fun, *primals, has_aux=False):
+def _vjp(fun, *primals, has_aux=False):
   canon = lambda x: x if isinstance(x, core.Tracer) else canonicalize_value(x)
   primals = tree_map(canon, primals)
   primals_flat, in_tree = tree_flatten(primals)

jax/_src/config.py

@@ -1833,12 +1833,6 @@ def _validate_default_device(val):
     help='Enable error checks for mutable arrays that rule out aliasing.',
     include_in_trace_context=True)
 
-vjp3 = bool_state(
-    name='jax_vjp3',
-    default=True,
-    upgrade=True,
-    help='Use new backward-pass code in jax.vjp')
-
 refs_to_pins = bool_state(
     name='jax_refs_to_pins',
     default=False,

jax/_src/interpreters/ad.py

@@ -26,7 +26,7 @@
 from jax._src import linear_util as lu
 from jax._src.interpreters import partial_eval as pe
 from jax._src.tree_util import (tree_flatten, tree_unflatten,
-                                register_pytree_node, Partial, PyTreeDef)
+                                register_pytree_node, PyTreeDef)
 from jax._src import mesh as mesh_lib
 from jax._src import core
 from jax._src import source_info_util
@@ -302,23 +302,6 @@ def linearize(traceable: lu.WrappedFun, *primals, **kwargs):
   else:
     return out_primals_consts, out_tangents_pvals, jaxpr, consts, aux()
 
-def vjp(traceable: lu.WrappedFun, primals, has_aux=False):
-  if not has_aux:
-    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
-  else:
-    out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)
-
-  def unbound_vjp(pvals, jaxpr, consts, *cts):
-    cts = tuple(ct for ct, pval in zip(cts, pvals) if not pval.is_known())
-    dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
-    arg_cts = backward_pass(jaxpr, True, consts, dummy_args, cts)
-    return map(instantiate_zeros, arg_cts)
-
-  vjp_ =  Partial(partial(unbound_vjp, pvals, jaxpr), consts)
-  if not has_aux:
-    return out_primals, vjp_
-  else:
-    return out_primals, vjp_, aux
 
 # NOTE: The FIXMEs below are caused by primal/tangent mixups (type
 # errors if you will)

jax/interpreters/ad.py

@@ -41,7 +41,6 @@
   primitive_jvps as primitive_jvps,
   primitive_transposes as primitive_transposes,
   reducing_transposes as reducing_transposes,
-  vjp as vjp,
   zeros_like_aval as zeros_like_aval,
 )
 

tests/api_test.py

@@ -3101,7 +3101,7 @@ def cond(pred):
   def test_grad_of_bool_vjp3(self):
     def cond(pred):
       return lax.cond(pred, lambda _: 1., lambda _: 2., 1.)
-    value, f_vjp = api.vjp3(cond, True)
+    value, f_vjp = api.vjp(cond, True)
     grd, = f_vjp(1.)
     self.assertEqual(value, 1.)
     self.assertEqual(grd, np.zeros(shape=(), dtype=float0))
@@ -3213,6 +3213,12 @@ def f():
       self.assertNotRegex(str(j_module),
        f"stablehlo.constant dense.*tensor<{const_size}x")
 
+  def test_basic_vjp3(self):
+    f = jax.jit(lambda x: jnp.sin(jnp.sin(x)))
+    _, f_vjp = jax.vjp(f, 1.)
+    g, = f_vjp(1.0)
+    self.assertAllClose(g, jnp.cos(jnp.sin(1.)) * jnp.cos(1.), check_dtypes=False)
+
   def test_constants_not_in_lowering_scan(self):
     if not config.use_simplified_jaxpr_constants.value:
       self.skipTest("Works only with simplified Jaxpr consts")
@@ -6737,11 +6743,7 @@ def f(x):
       return lax.cond(x.sum() > 0, f, lambda x: x, x)
 
     _, f_vjp = api.vjp(f, jnp.ones((5, 5)))
-    if config.vjp3.value:
-      jaxpr_text = str(f_vjp.jaxpr)
-    else:
-      jaxpr_text = str(f_vjp.jaxpr)
-
+    jaxpr_text = str(f_vjp.jaxpr)
     self.assertEqual(jaxpr_text.count(' sin '), 2)
     self.assertEqual(jaxpr_text.count(' cos '), 3)
     # Five calls to dot_general in the backward pass because we have two for
@@ -7806,7 +7808,7 @@ def test_basic_unused(self):
   def test_basic_unused_vjp3(self):
     f = jnp.sin
     primals = 3.,
-    y, f_vjp = api.vjp3(f, *primals)
+    y, f_vjp = api.vjp(f, *primals)
     x_ct, = f_vjp(1.)
     self.assertAllClose(y, jnp.sin(3.))
     self.assertAllClose(x_ct, jnp.cos(3.))
@@ -7821,7 +7823,7 @@ def test_basic_opaque(self):
   def test_basic_opaque_vjp3(self):
     f = jnp.sin
     primals = 3.,
-    _, f_vjp = api.vjp3(f, *primals)
+    _, f_vjp = api.vjp(f, *primals)
     assert f_vjp.opaque_residuals  # can detect if opaque res are used
 
   def test_basic_pytree_error(self):
@@ -7841,7 +7843,7 @@ def f(x):
   #   def f(x):
   #     return [x['hi'] * x['bye']]
 
-  #   y, f_vjp = api.vjp3(f, {'hi': 2., 'bye': 3.})
+  #   y, f_vjp = api.vjp(f, {'hi': 2., 'bye': 3.})
   #   arg_ct, = f_vjp([1.], {'hi': 2., 'bye': 3.})
   #   self.assertAllClose(y, [6.])
   #   self.assertAllClose(arg_ct, {'hi': 3., 'bye': 2.})
@@ -7890,7 +7892,7 @@ def f2(x, w):
 
     x = jnp.ones((3, 4))
     w = jnp.ones((4, 4))
-    y, f2_vjp = api.vjp3(f2, x, w)
+    y, f2_vjp = api.vjp(f2, x, w)
     f2_vjp.args_res[1] = None
     y_grad = jnp.ones_like(y)
     f2_vjp.args_res[1] = w
@@ -7964,7 +7966,7 @@ def foo(x):
 
   def test_grad_traceback(self):
     # TODO(dougalm): improve this
-    expected_depth = 12
+    expected_depth = 11
     init_depth = self.cur_depth()
 
     def foo(x):
@@ -7987,7 +7989,7 @@ def foo(x):
   def test_custom_vjp_traceback(self):
     # TODO(dougalm): improve this
     expected_depth_f = 10
-    expected_depth_f_fwd = 20
+    expected_depth_f_fwd = 19
     expected_depth_f_rev = 12
     init_depth = self.cur_depth()
     @jax.custom_vjp

tests/custom_api_test.py

@@ -1472,7 +1472,7 @@ def sin_jvp(primals, tangents):
       (x, y), (x_dot, y_dot) = primals, tangents
       del y_dot  # ignore lol
       return div(x, y), div(x_dot, y)
-    _, f_vjp = api.vjp3(lambda x: div(x, 2.), 1.)
+    _, f_vjp = api.vjp(lambda x: div(x, 2.), 1.)
     ans, = f_vjp(1.)
     self.assertAllClose(ans, 1./2, check_dtypes=False)
 

tests/lax_control_flow_test.py

@@ -41,7 +41,6 @@
 import jax.numpy as jnp  # scan tests use numpy
 import jax.scipy as jsp
 from jax._src import dispatch
-from jax._src.api import vjp3
 from jax._src.lax import control_flow as lax_control_flow
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
@@ -2777,10 +2776,7 @@ def cumprod(x):
     # ==> Yes, we don't want to change autodiff const behavior. We must make
     # these tessts pass under use_simplified_jaxpr_constants.
     if not config.use_simplified_jaxpr_constants.value:
-      if config.vjp3.value:
-        ext_res, = vjp_fun.args_res
-      else:
-        *_, ext_res = vjp_fun.args[0].args[0]
+      ext_res, = vjp_fun.args_res
       self.assertIs(ext_res, x)
 
     if remat is not None:
@@ -2790,10 +2786,7 @@ def cumprod(x):
     x = rng.randn(32, 2, 32).astype('float32')  # numpy.ndarray, not Array
     _, vjp_fun = jax.vjp(cumprod, x)
     if not config.use_simplified_jaxpr_constants.value:
-      if config.vjp3.value:
-        ext_res, *_ = vjp_fun.opaque_residuals
-      else:
-        *_, ext_res = vjp_fun.args[0].args[0]
+      ext_res, *_ = vjp_fun.opaque_residuals
       self.assertIsInstance(ext_res, jax.Array)
 
   def test_scan_vmap_collectives(self):
@@ -3498,14 +3491,14 @@ def test_cond_basic_vjp3(self):
     def f(x):
       return jax.lax.cond(True, jnp.sin, lambda x: x, x)
 
-    _, f_vjp = vjp3(f, 1.)
+    _, f_vjp = jax.vjp(f, 1.)
     g, = f_vjp(1.0)
     self.assertAllClose(g, jnp.cos(1.), check_dtypes=False)
 
     def h(x):
       return jax.lax.cond(True, jnp.sin, lambda x: 1., x)
 
-    _, h_vjp = vjp3(h, 1.)
+    _, h_vjp = jax.vjp(h, 1.)
     g, = h_vjp(1.0)
     self.assertAllClose(g, jnp.cos(1.), check_dtypes=False)
 

tests/mutable_array_test.py

@@ -25,7 +25,6 @@
 from jax._src import core
 from jax._src import config
 from jax._src import test_util as jtu
-from jax._src.api import vjp3
 from jax._src.util import safe_map, safe_zip
 from jax._src.interpreters import mlir
 from jax.sharding import NamedSharding, PartitionSpec as P, AxisType
@@ -524,8 +523,8 @@ def stash_grads_bwd(grads_ref, g):
 
     grads_ref = core.new_ref(jnp.float32(0.))
     x = jnp.float32(1.)
-    _, f_vjp, *maybe_aux = vjp3(lambda x: primal(grads_ref, x), x,
-                               has_aux=has_aux)
+    _, f_vjp, *maybe_aux = jax.vjp(
+        lambda x: primal(grads_ref, x), x, has_aux=has_aux)
     _ = f_vjp(jnp.float32(1.))
     self.assertAllClose(grads_ref[...], jnp.cos(jnp.sin(1.)), check_dtypes=False)
     if has_aux:
@@ -553,15 +552,15 @@ def stash_grads_bwd(stash_ref, g):
     stash_grads.defvjp(stash_grads_fwd, stash_grads_bwd)
 
     stash_ref = core.new_ref(jnp.float32(0.))
-    _, f_vjp = vjp3(lambda x: primal(stash_ref, x), jnp.float32(1.))
+    _, f_vjp = jax.vjp(lambda x: primal(stash_ref, x), jnp.float32(1.))
     grads_val, = f_vjp(jnp.float32(1.))
     self.assertAllClose(stash_ref[...], jnp.cos(jnp.sin(1.)), check_dtypes=False)
     self.assertAllClose(grads_val, jnp.cos(jnp.sin(1.)) * jnp.cos(1.),
                         check_dtypes=False)
 
     stash_ref = core.new_ref(jnp.float32(0.))
     grads_ref = core.new_ref(jnp.float32(0.))
-    _, f_vjp = vjp3(lambda x: primal(stash_ref, x), jnp.float32(1.))
+    _, f_vjp = jax.vjp(lambda x: primal(stash_ref, x), jnp.float32(1.))
     _ = f_vjp.with_refs(grads_ref)(jnp.float32(1.))
     self.assertAllClose(stash_ref[...], jnp.cos(jnp.sin(1.)), check_dtypes=False)
     self.assertAllClose(grads_ref[...], jnp.cos(jnp.sin(1.)) * jnp.cos(1.),
@@ -849,7 +848,7 @@ def process_batch(Ws, xs_batch):
       grad_acc = jax.new_ref(jnp.zeros_like(Ws))               # CHANGED
 
       def process_mubatch(_, xs):
-        loss, f_vjp = vjp3(lambda Ws: mubatch_loss(Ws, xs), Ws)  # CHANGED
+        loss, f_vjp = jax.vjp(lambda Ws: mubatch_loss(Ws, xs), Ws)  # CHANGED
         f_vjp.with_refs(grad_acc)(jnp.ones_like(loss))           # CHANGED
         return (), loss
 
@@ -924,7 +923,7 @@ def f_bwd(_, g):
     self.assertAllClose(y, 3.14, check_dtypes=False)
 
     # this exercises the fallback path, not a fancy transpose
-    _, f_vjp = vjp3(lambda x: f(jax.new_ref(x)), 3.14)
+    _, f_vjp = jax.vjp(lambda x: f(jax.new_ref(x)), 3.14)
     g, = f_vjp(1.)
     self.assertAllClose(g, 1., check_dtypes=False)
 
@@ -969,7 +968,7 @@ def body(_, xy):
       return z.sum()
 
     grad_accum = jax.new_ref(jnp.zeros(5))
-    _, f_vjp = vjp3(f, jnp.ones(5))
+    _, f_vjp = jax.vjp(f, jnp.ones(5))
     _, = f_vjp.with_refs(grad_accum)(1.)
     self.assertAllClose(grad_accum[...], jnp.arange(5.))
 
@@ -978,7 +977,7 @@ def test_vmap_with_vjp3(self):
     def grad_via_ref(f):
       def wrapper(*args):
         grad_accum = jax.tree.map(lambda x: jax.new_ref(jnp.zeros_like(x)), args)
-        out, f_vjp = vjp3(f, *args)
+        out, f_vjp = jax.vjp(f, *args)
         f_vjp.with_refs(*grad_accum)(jnp.ones_like(out))
         return jax.tree.map(lambda x: jax.freeze(x), grad_accum)
       return wrapper

tests/pjit_test.py

@@ -38,7 +38,6 @@
 from jax._src import dtypes
 from jax import stages
 from jax import lax
-from jax._src.api import vjp3
 from jax._src.lax import lax as lax_internal
 from jax.lax import with_sharding_constraint
 from jax._src import prng
@@ -1309,12 +1308,6 @@ def test_device_put_copy_donate(self):
     self.assertNotDeleted(z)
     self.assertArraysEqual(a, x * 2)
 
-  def test_basic_vjp3(self):
-    f = jax.jit(lambda x: jnp.sin(jnp.sin(x)))
-    _, f_vjp = vjp3(f, 1.)
-    g, = f_vjp(1.0)
-    self.assertAllClose(g, jnp.cos(jnp.sin(1.)) * jnp.cos(1.), check_dtypes=False)
-
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class AutoShardingPjitTest(jtu.JaxTestCase):