Merge pull request #32735 from mattjj:vjp3-args-res-tuple-trees

PiperOrigin-RevId: 820932861

Author: Google-ML-Automation

jax/_src/api.py

@@ -2292,7 +2292,8 @@ def _vjp3(fun, *primals, has_aux=False):
   spec = [used.add(id(r)) or RSpec(id_map[id(r)], True) if id(r) in id_map else  # type: ignore
           RSpec(opaque_residuals.append(r) or (len(opaque_residuals) - 1), False)  # type: ignore
           for r in residuals]
-  args_res = tree_map(lambda x: x if id(x) in used else NotNeeded(), primals)
+  args_res = tuptree_map(lambda x: x if id(x) in used else NotNeeded(),
+                         in_tree, primals)
   out_primal_avals = [typeof(x) for x in out_primals_flat]
   f_vjp = VJP(partial(_vjp3_callable, spec, out_known, jaxpr, out_primal_avals),
               in_tree, out_tree, list(args_res), opaque_residuals)
@@ -2302,6 +2303,9 @@ def _vjp3(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)
@@ -2311,10 +2315,8 @@ def _vjp3_callable(spec, out_known, jaxpr, out_primal_avals, in_tree, out_tree,
                    args_res, opaque_res, *maybe_ct_refs):
   maybe_ct_refs_flat, in_tree_ = tree_flatten(maybe_ct_refs)
   if in_tree != in_tree_: raise Exception
-  args_res_flat, in_tree_ = tree_flatten(
-      tuple(args_res), is_leaf=lambda x: isinstance(x, NotNeeded))
-  if in_tree != in_tree_: raise Exception
-  residuals = [args_res_flat[i.idx] if i.primal else opaque_res[i.idx] for i in spec]
+  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)
                 for v, x in zip(jaxpr.invars, maybe_ct_refs_flat)]
   return Partial(partial(_vjp3_bwd, in_tree, out_tree, out_known, jaxpr,

tests/api_test.py

@@ -7817,12 +7817,27 @@ def test_basic_unused(self):
     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.,
+    y, f_vjp = api.vjp3(f, *primals)
+    x_ct, = f_vjp(1.)
+    self.assertAllClose(y, jnp.sin(3.))
+    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.vjp3(f, *primals)
+    assert 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']]
@@ -7835,6 +7850,20 @@ def f(x):
     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.vjp3(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):
@@ -7849,6 +7878,24 @@ def f2(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)
+
+  def test_fsdp_vjp3(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_vjp = api.vjp3(f2, x, w)
+    f2_vjp.args_res[1] = None
+    y_grad = jnp.ones_like(y)
+    f2_vjp.args_res[1] = w
+    x_grad, w_grad = f2_vjp(y_grad)
+    self.assertAllClose(x_grad, 2. * y_grad @ w.T)
+    self.assertAllClose(w_grad, 2. * x.T @ y_grad)
     self.assertAllClose(w_grad, 2. * x.T @ y_grad)
 
   def test_doesnt_leak_symbolic_zeros(self):