Delete si_vjp from JAX. Dead weight at this point. It is now replaced with jax.vjp.

PiperOrigin-RevId: 844865984

Author: yashk2810

CHANGELOG.md

@@ -37,6 +37,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     For the moment, during Python type checking, we continue to declare `Tracer`
     as a subclass of `Array`, however we expect to remove this in a future
     release.
+  * `jax.experimental.si_vjp` has been deleted.
+    `jax.vjp` subsumes it's functionality.
 
 ## JAX 0.8.1 (November 18, 2025)
 

jax/_src/api.py

@@ -2212,106 +2212,12 @@ def vjp(
       fun, debug_info=debug_info("vjp", fun, primals, {}))
   return _vjp(wrapped_fun, *primals, has_aux=has_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):
-  if len(which) != len(primals):
-    raise ValueError(
-    "length of 'which' argument must equal the number of primal input values, "
-    f"but got {len(which)=} and {len(primals)=}")
-
-  dbg = debug_info("saved_input_vjp", f, primals, {})
-  fun = lu.wrap_init(f, debug_info=dbg)
-  primals_flat, in_tree = tree_flatten(primals)
-  fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-  out_primals_flat, out_pvals, jaxpr, residuals = ad.linearize(fun, *primals_flat)
-  out_known = [pval.is_known() for pval in out_pvals]
-  primals_filt, filt_tree = tree_flatten(tuple(p for w, p in zip(which, primals) if w))
-  id_map = {id(x): i for i, x in enumerate(primals_filt)}
-  opaque_residuals = []
-  res_spec = [RSpec(id_map[id(r)], True) if id(r) in id_map else
-              RSpec(opaque_residuals.append(r) or (len(opaque_residuals) - 1), False)  # type: ignore
-              for r in residuals]
-  out_primal_avals = map(shaped_abstractify, out_primals_flat)
-  f_vjp = Partial(partial(_saved_input_vjpfun, res_spec, filt_tree, in_tree,
-                          out_tree(), out_known, jaxpr, out_primal_avals),
-                  opaque_residuals)
-
-  if not allow_unused and not set(id_map).issubset(res_ids := {id(r) for r in residuals}):
-    unused = [(i, core.get_aval(x)) for i, (x, w) in enumerate(zip(primals, which))
-              if w and id(x) not in res_ids]
-    assert unused
-    if len(unused) == 1:
-      (i, a), = unused
-      start, was = "an input value", "was"
-      msg = f" {dbg.arg_names[i] if dbg.arg_names is not None else 'unknown'} of type {a.str_short()}"
-    else:
-      start, was = "multiple input values", "were"
-      msg = "\n" + "\n".join(f"  * {dbg.arg_names[i] if dbg.arg_names is not None else 'unknown'} of type {a.str_short()}"
-                             for i, a in unused)
-    raise Exception(f"with {allow_unused=}, {start} marked to be saved {was} "
-                    f"not used by the backward pass:{msg}")
-
-  if not allow_opaque and opaque_residuals:
-    msg = ", ".join(core.get_aval(x).str_short() for x in opaque_residuals)
-    raise Exception(f"with {allow_opaque=}, the backward pass requires opaque "
-                    f"(non-input) residuals: {msg}")
-
-  out_primals = tree_unflatten(out_tree(), out_primals_flat)
-  return out_primals, f_vjp
-
-def _saved_input_vjpfun(res_spec, filtered_tree, in_tree, out_tree, out_known,
-                        jaxpr, out_primal_avals, opaque_residuals, ct,
-                        *saved_primals):
-  primals_filtered, filtered_tree_ = tree_flatten(saved_primals)
-  if filtered_tree != filtered_tree_:
-    raise ValueError(
-        "inputs passed to f_vjp must be a tuple of (pytrees of) "
-        "arrays with the same structure as\n"
-        "  tuple(x for x, w in zip(inputs, which) if w)\n"
-        "given the original call\n"
-        "  _, f_vjp = saved_input_vjp(f, which, *inputs, ...)\n"
-        "but the structures differ:\n" +
-        "\n".join(f"  * inputs{keystr(path)} was a {thing1} in the original "
-                  f"call, but a {thing2} here, so {explanation}"
-                  for path, thing1, thing2, explanation
-                  in equality_errors_pytreedef(filtered_tree, filtered_tree_)))
-
-  residuals = [primals_filtered[i.idx] if i.primal else opaque_residuals[i.idx]
-               for i in res_spec]
-  dummy_args = [ad.UndefinedPrimal(v.aval) for v in jaxpr.invars]
-  cts_flat, out_tree_ = tree_flatten(ct)
-  if out_tree_ != out_tree:
-    raise ValueError(f"unexpected tree structure of argument to vjp function: "
-                     f"got {out_tree_}, but expected to match {out_tree}")
-  for arg, aval in zip(cts_flat, out_primal_avals):
-    ct_aval = shaped_abstractify(arg)
-    ct_aval_expected = aval.to_cotangent_aval()
-    if (not core.typecompat(ct_aval, ct_aval_expected) and
-        not _temporary_dtype_exception(ct_aval, ct_aval_expected)):
-      raise ValueError(
-          "unexpected JAX type (e.g. shape/dtype) for argument to vjp function: "
-          f"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} "
-          f"because the corresponding output of the function had JAX type "
-          f"{aval.str_short()}")
-
-  cts_flat = [ct for ct, k in zip(cts_flat, out_known) if not k]
-  arg_cts = ad.backward_pass(jaxpr, True, residuals, dummy_args, cts_flat)
-  return tree_unflatten(in_tree, map(ad.instantiate_zeros, arg_cts))
-
-@dataclasses.dataclass(frozen=True)
-class RSpec:
-  idx: int
-  primal: bool
-
-si_vjp = saved_input_vjp
-
-
 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)
-  for arg in primals_flat: dispatch.check_arg(arg)
+  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_flat, out_pvals, jaxpr, residuals = ad.linearize(
@@ -2340,22 +2246,14 @@ def _vjp(fun, *primals, has_aux=False):
   else:
     return out_primals, f_vjp, tree_unflatten(aux_tree, aux)
 
-def tuptree_map(f, treedef, x):
-  return treedef.walk(lambda xs, _: tuple(xs), f, x)
-
-
-def _is_ref(x):
-  from jax._src.state.types import AbstractRef
-  try: return isinstance(typeof(x), AbstractRef)
-  except: return False
-
 def _vjp3_callable(spec, out_known, jaxpr, out_primal_avals, in_tree, out_tree,
                    args_res, opaque_res, *maybe_ct_refs):
   if not maybe_ct_refs:
     maybe_ct_refs_flat = [GradValue()] * in_tree.num_leaves
   else:
     maybe_ct_refs_flat, in_tree_ = tree_flatten(maybe_ct_refs)
-    if in_tree != in_tree_: raise Exception  # TODO accept isomorph tuple tree
+    if in_tree != in_tree_:
+      raise Exception  # TODO accept isomorph tuple tree
   args_res_ = tree_leaves(args_res, is_leaf=lambda x: isinstance(x, NotNeeded))
   residuals = [args_res_[i.idx] if i.primal else opaque_res[i.idx] for i in spec]
   maybe_refs = [ad.RefAccum(v.aval, x) if _is_ref(x) else ad.ValAccum(v.aval)
@@ -2366,7 +2264,8 @@ def _vjp3_callable(spec, out_known, jaxpr, out_primal_avals, in_tree, out_tree,
 def _vjp3_bwd(in_tree, out_tree, out_known, jaxpr, out_primal_avals, residuals,
               maybe_refs, out_ct):
   cts_flat, out_tree_ = tree_flatten(out_ct)
-  if out_tree != out_tree_: _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)
+  if out_tree != out_tree_:
+    _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)
   _vjp_check_ct_avals(cts_flat, out_primal_avals)
   cts_flat = [ct for ct, k in zip(cts_flat, out_known) if not k]
   ad.backward_pass3(jaxpr, True, residuals, maybe_refs, cts_flat)
@@ -2375,6 +2274,23 @@ def _vjp3_bwd(in_tree, out_tree, out_known, jaxpr, out_primal_avals, residuals,
   arg_cts = map(ad.instantiate_zeros, arg_cts)
   return tree_unflatten(in_tree, arg_cts)
 
+
+@dataclasses.dataclass(frozen=True)
+class RSpec:
+  idx: int
+  primal: bool
+
+def tuptree_map(f, treedef, x):
+  return treedef.walk(lambda xs, _: tuple(xs), f, x)
+
+def _is_ref(x):
+  from jax._src.state.types import AbstractRef
+  try:
+    return isinstance(typeof(x), AbstractRef)
+  except:
+    return False
+
+
 _vjp_too_many_args = """
 The function returned by `jax.vjp` applied to {} was called with {} arguments,
 but functions returned by `jax.vjp` must be called with a single argument
@@ -2396,6 +2312,7 @@ def f(x):
 arguments rather than in a tuple, this error can arise.
 """.format
 
+
 def _vjp_ct_tree_error(jaxpr, out_tree, ct_tree):
   msg = f"""unexpected tree structure.
 
@@ -2410,6 +2327,7 @@ def _vjp_ct_tree_error(jaxpr, out_tree, ct_tree):
                    in equality_errors_pytreedef(out_tree, ct_tree))
   raise ValueError(msg)
 
+
 def _vjp_check_ct_avals(cts, primal_avals):
   # TODO(mattjj): improve this error  by flattening with keys in the first place
   for ct, aval in zip(cts, primal_avals):
@@ -2425,6 +2343,7 @@ def _vjp_check_ct_avals(cts, primal_avals):
           "because the corresponding output of the differentiated function had JAX type "
           f"{aval.str_short()}")
 
+
 @register_dataclass
 @dataclasses.dataclass(frozen=True)
 class NotNeeded:

jax/experimental/__init__.py

@@ -21,10 +21,6 @@
 # experimental features and as a result, more flexibility to manage their status
 # and lifetimes.
 
-from jax._src.api import (
-  saved_input_vjp as saved_input_vjp,
-  si_vjp as si_vjp
-)
 from jax._src.callback import (
   io_callback as io_callback
 )

tests/api_test.py

@@ -7769,9 +7769,12 @@ def test_basic(self):
     def f(x, y):
       return x * y
 
-    primals = 2., 3.
-    y, f_vjp = api.si_vjp(f, [True, True], *primals)
-    arg_cts = f_vjp(1., *primals)
+    primals = [2., 3.]
+    y, f_vjp = jax.vjp(f, *primals)
+    f_vjp.args_res = [None, None]
+    y_grad = 1.
+    f_vjp.args_res = primals
+    arg_cts = f_vjp(1.)
     self.assertAllClose(y, 6.)
     self.assertAllClose(arg_cts, (3., 2.))
 
@@ -7782,29 +7785,20 @@ def f(x, y):
     @jax.jit
     def g():
       primals = 2., 3.
-      y, f_vjp = api.si_vjp(f, [True, True], *primals)
+      y, f_vjp = jax.vjp(f, *primals)
+      f_vjp.args_res = [None, None]
       return y, f_vjp
 
     @jax.jit
     def h(f_vjp):
-      return f_vjp(1., 2., 3.)
+      f_vjp.args_res = [2., 3.]
+      return f_vjp(1.)
 
     y, f_vjp = g()
     arg_cts = h(f_vjp)
     self.assertAllClose(y, 6.)
     self.assertAllClose(arg_cts, (3., 2.))
 
-  def test_basic_unused(self):
-    f = jnp.sin
-    primals = 3.,
-    y, f_vjp = api.si_vjp(f, [True], *primals)
-    x_ct, = f_vjp(1., *primals)
-    self.assertAllClose(y, jnp.sin(3.))
-    self.assertAllClose(x_ct, jnp.cos(3.))
-
-    with self.assertRaisesRegex(Exception, "not used by the backward pass: x"):
-      _ = api.si_vjp(f, [True], *primals, allow_unused=False)
-
   def test_basic_unused_vjp3(self):
     f = jnp.sin
     primals = 3.,
@@ -7814,58 +7808,28 @@ def test_basic_unused_vjp3(self):
     self.assertAllClose(x_ct, jnp.cos(3.))
     self.assertIsInstance(f_vjp.args_res[0], api.NotNeeded)  # can check if unused
 
-  def test_basic_opaque(self):
-    f = jnp.sin
-    primals = 3.,
-    with self.assertRaisesRegex(Exception, "the backward pass requires opaque"):
-      _ = api.si_vjp(f, [True], *primals, allow_opaque=False)
-
   def test_basic_opaque_vjp3(self):
     f = jnp.sin
     primals = 3.,
     _, f_vjp = api.vjp(f, *primals)
-    assert f_vjp.opaque_residuals  # can detect if opaque res are used
+    self.assertTrue(f_vjp.opaque_residuals)  # can detect if opaque res are used
 
   def test_basic_pytree_error(self):
     def f(x):
       return [x['hi'] * x['bye']]
 
-    y, f_vjp = api.si_vjp(f, [True], {'hi': 2., 'bye': 3.})
-    arg_ct, = f_vjp([1.], {'hi': 2., 'bye': 3.})
+    y, f_vjp = jax.vjp(f, {'hi': 2., 'bye': 3.})
+    f_vjp.args_res = [None]
+    y_grad = [1.]
+    f_vjp.args_res = [{'hi': 2., 'bye': 3.}]
+    arg_ct, = f_vjp(y_grad)
     self.assertAllClose(y, [6.])
     self.assertAllClose(arg_ct, {'hi': 3., 'bye': 2.})
 
-    with self.assertRaisesRegex(ValueError, "but the structures differ"):
-      f_vjp(1., {'hi': 2.})
-
-  # TODO(mattjj): improve this vjp3 error message
-  # def test_basic_pytree_error_vjp3(self):
-  #   def f(x):
-  #     return [x['hi'] * x['bye']]
-
-  #   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.})
-
-  #   f_vjp.args_res[0] = {'hi': 2.}
-  #   with self.assertRaisesRegex(ValueError, "but the structures differ"):
-  #     f_vjp(1.)
-
-  def test_fsdp(self):
-    # see https://github.com/jax-ml/jax/pull/27017 for why this is called "fsdp"
-    def f2(x, w):
-      x = 1. * x
-      x = x @ w
-      x = 2. * x
-      return x
-
-    x = jnp.ones((3, 4))
-    w = jnp.ones((4, 4))
-    y, f2_sivjp = api.si_vjp(f2, [False, True], x, w)
-    y_grad = jnp.ones_like(y)
-    x_grad, w_grad = f2_sivjp(y_grad, w)
-    self.assertAllClose(x_grad, 2. * y_grad @ w.T)
+    # TODO(mattjj): Raise an error message.
+    # with self.assertRaisesRegex(ValueError, "but the structures differ"):
+    #   f_vjp.args_res = [{'hi': 2.}]
+    #   f_vjp([1.])
 
   def test_fsdp_error(self):
     # see https://github.com/jax-ml/jax/pull/27017 for why this is called "fsdp"
@@ -7877,10 +7841,12 @@ def f2(x, w):
 
     x = jnp.ones((3, 4))
     w = jnp.ones((4, 4))
-    y, f2_sivjp = api.si_vjp(f2, [False, True], x, w)
+    y, f2_vjp = jax.vjp(f2, x, w)
+    f2_vjp.args_res[1] = None
     y_grad = jnp.ones((2, 4))
+    f2_vjp.args_res[1] = w
     with self.assertRaisesRegex(ValueError, "unexpected JAX type"):
-      f2_sivjp(y_grad, w)
+      f2_vjp(y_grad)
 
   def test_fsdp_vjp3(self):
     # see https://github.com/jax-ml/jax/pull/27017 for why this is called "fsdp"
@@ -7902,10 +7868,11 @@ def f2(x, w):
     self.assertAllClose(w_grad, 2. * x.T @ y_grad)
 
   def test_doesnt_leak_symbolic_zeros(self):
-    _, vjp = api.si_vjp(lambda x: 1., [False], 3.14)
+    _, vjp = jax.vjp(lambda x: 1., 3.14)
     ans, = vjp(1.0)
     self.assertIsInstance(ans, jax.Array)
 
+
 class TracebackTest(jtu.JaxTestCase):
   # These tests are to catch regressions in Python traceback sizes. Our
   # second-order APIs can be nested arbitrarily and if each one adds a dozen