Repository: google/jax

Issue:
How do you nest implementations of VJP?
If I have a function `f` that has a VJP.  Can I somehow build a VJP on top of it for some function `g`?  My use case is that the backwards pass of `g` is quite complicated and makes multiple calls to `f_vjp`.   Here's a simplified example:
```python
from jax import vjp, custom_vjp

def f(x):
    return x ** 2

@custom_vjp
def g(x):
    return f(x)

def g_fwd(x):
    return vjp(f, x)

def g_bwd(f_vjp, y_bar):
    return f_vjp(y_bar)

g.defvjp(g_fwd, g_bwd)
y, g_vjp = vjp(g, 1.0)
```
This prints: `TypeError: <class 'functools.partial'> is not a valid JAX type` because `vjp` returns a `functool.partial` instance instead of a pytree.  It should be possible for `vjp` to return a pytree-like callable instead since internally `vjp(f, x)` produces residuals that must be pytree-like.  It's just unfortunate that when they're wrapped up into a callable, then that callable is not pytree-like.


Discussion:
I've looked at it a bit, and it appears that the issue is that jax.api._vjp returns
```
  vjp_py = partial(_vjp_pullback_wrapper, out_vjp,
                   [_dtype(x) for x in out_primal], (out_tree, in_tree))
```
where `out_vjp` is
```
    out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
```
and in `ad.vjp`, we see that `out_vjp` has some closed over values that are JAX trees:
```python
  out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  def vjp_(*cts):
    cts = tuple(map(ignore_consts, cts, pvals))
    dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
    arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
    return map(instantiate_zeros, arg_cts)
  return out_primals, vjp_
```
I guess `pvals`, `jaxprs`, and `consts` can all contain PyTrees?  Could some of them be static instead?  If they're definitely PyTrees, I might be able to make this change.  Or is there a better way to nest VJPs that I'm missing?
This seems like a case where perhaps you are pushing the limits of what `custom_vjp` is designed to support should be considering other options :)

That said, one way to fix this immediate issue would be to replace `partial` inside `jax.api._vjp` with `tree_util.Partial`, which is serializable as a pytree. You could also do this wrapping in user code:
```python
from jax import vjp, custom_vjp, tree_util

def f(x):
    return x ** 2

@custom_vjp
def g(x):
    return f(x)

def g_fwd(x):
    y, f_vjp = vjp(f, x)
    return y, tree_util.Partial(f_vjp)

def g_bwd(f_vjp, y_bar):
    return f_vjp(y_bar)

g.defvjp(g_fwd, g_bwd)
y, g_vjp = vjp(g, 1.0)
print(y, g_vjp(1.0))
# 1.0 (DeviceArray(2., dtype=float32),)
```

In general, it's fine to using `Partial` inside the forward pass of `custom_vjp` functions as long as you are careful not to close over any tracers, like the value `x` in this case. The dtypes, constants, jaxprs and treedefs used in the closures should all be fine. (Otherwise you should get a nasty error message.)

(@mattjj please correct me if I'm mis-stating anything here)
@shoyer Nice!!!  Brilliant solution.

As for other options, I think this is the simplest design: All of the nodes in my model have a regular VJP.  It's just that one of the node types' VJP would normally be very complicated, so I implement it as a custom VJP, which internally delegates some of the work to a nested VJP.  It seems to me to be very logical.

I find nesting patterns to be very logical:  a function delegating to a nested function, a list display delegating to a nested list display, etc.  My (possibly naive) feeling is that VJP and JVP should also be able to delegate to nested instances of themselves.

Since `vjp` already ensures that its arguments are pytrees, do you think it would be possible to make `vjp` return a `tree_util.Partial`?
> Since `vjp` already ensures that its arguments are pytrees, do you think it would be possible to make `vjp` return a `tree_util.Partial`?

@mattjj is the real expert here, so I'll defer to his judgment. But this seems reasonable to me.

@shoyer I tried to apply your solution to my code, but when used in a more complicated setting, I'm getting leaked tracers.  Why did you say:

> as long as you are careful not to close over any tracers, like the value x in this case.

Where is x closed over?

Can you think of any alternative way of nesting VJPs?  Am I missing something?

I think the reason I'm getting leaked tracers is because  `tree_view.Partial` assumes that the callable is necessarily static (not always true), and more importantly it assumes that all of the arguments are pytrees (very often false).  In particular, some of the itnernals of vjp are passing static arguments as positional arguments to partial (lists of dtypes, other functions, lists of PyTreeDef).

I think the solution is to write a more advanced `Partial` that marks its static parameters.  Then, this can be used in the VJP codepath in JAX's codebase.  I don't mind doing it if you think it's a good idea.  What do you think?
How are you calling `Partial`? On further consideration, maybe we don't want to substitute `partial` -> `Partial`, but rather just add `Partial` on top as a wrapper.

In my example, `Partial` has no explicit argument applied at all, and is just a way of marking the function as a pytree. That might be a better approach.
Thinking about this a little more carefully: of course you need to explicitly wrap values that might include tracers.

`pvals` and `consts` are the parameters that might include tracers. Both should already be pytrees.
@shoyer Thanks a lot for exploring this.  I'm just passing the result of `vjp` into `Partial` just like in your comment above.  The reason this doesn't work is because the function outputted by `vjp` is a `partial` whose arguments include static arguments.

Your code above does this:
```
f_vjp: partial(...)
g = Partial(f_vjp)
```
Here, `g` is a pytree whose only component is a static `partial` object.  But that partial object contains within it some pytrees, so it leaks tracers.

If you try to do this:
```
g = Partial(f_vjp.func, *f_vjp.args)
```
then you find out that `f_vjp.args` also contains static components, and JAX complains that these components are not pytrees.

I guess pvals and consts are not problematic.  Things which are problematic are the function stored as an argument (I think it was `vjp_`), and a list of dtypes, both of which need to be static.

I think the long story short of this is that we might consider changing JAX's VJP codepath so that the partial function applications that it creates are created with `Partial`, and to ensure that its arguments are all pytrees.  It might be convenient to create a simple box (using `flax.struct.dataclass`) like
```python
@dataclass
class StaticBox:
    value = field(False)
```
or equivalently by writing out the flatten and unflatten methods.  You could then put the static parameters in boxes, and send them to Partial, and it should all work.
I'm trying to use this version of `Partial` to ensure that the output of `vjp` doesn't leak tracers:
```python
class Partial(functools.partial):
  """A version of functools.partial that works in pytrees.

  Use it for partial function evaluation in a way that is compatible with JAX's
  transformations, e.g., ``Partial(func, *args, **kwargs)``.

  (You need to explicitly opt-in to this behavior because we didn't want to give
  functools.partial different semantics than normal function closures.)
  """
  def __new__(cls, func, /, *args,
              callable_is_static=True,
              static_argnums=(),
              static_kwargs={},
              **kwargs):
    if isinstance(func, Partial):
      raise TypeError
    retval = super().__new__(cls, func, *args, **kwargs)
    retval.callable_is_static = callable_is_static
    retval.static_argnums = set(static_argnums)
    retval.static_kwargs = static_kwargs
    return retval

  def tree_flatten(self):
    static_args = []
    tree_args = []

    def _append(is_static, value):
      if is_static:
        static_args.append(value)
      else:
        tree_args.append(value)

    _append(self.callable_is_static, self.func)
    for i, value in enumerate(self.args):
      _append(i in self.static_argnums, value)

    return ((list(reversed(tree_args)), self.keywords),
            (self.callable_is_static, self.static_argnums,
             list(reversed(static_args)), self.static_kwargs))

  @classmethod
  def tree_unflatten(cls, static, trees):
    callable_is_static, static_argnums, static_args, static_kwargs = static
    tree_args, tree_kwargs = trees
    args = []

    for i in range(len(static_args) + len(tree_args)):
      if i == 0:
        is_static = callable_is_static
      else:
        is_static = i - 1 in static_argnums
      if is_static:
        args.append(static_args.pop(0))
      else:
        args.append(tree_args.pop(0))

    return Partial(*reversed(args),
                   callable_is_static=callable_is_static,
                   static_argnums=static_argnums,
                   static_kwargs=static_kwargs,
                   **tree_kwargs)

  def __call__(self, *args, **kwargs):
    super().__call__(*args, **self.static_kwargs, **kwargs)
```
Still getting leaked tracers, but it's getting farther than before.
I don't think this needs a new Pytree. You could just make a function that
uses a combination of functools.partial and tree_util.Partial.

On Thu, Jul 9, 2020 at 9:57 AM Neil <notifications@github.com> wrote:

> I'm trying to use this version of Partial to ensure that the output of vjp
> doesn't leak tracers:
>
> class Partial(functools.partial):
>   """A version of functools.partial that works in pytrees.  Use it for partial function evaluation in a way that is compatible with JAX's  transformations, e.g., ``Partial(func, *args, **kwargs)``.  (You need to explicitly opt-in to this behavior because we didn't want to give  functools.partial different semantics than normal function closures.)  """
>   def __new__(cls, func, /, *args,
>               callable_is_static=True,
>               static_argnums=(),
>               static_kwargs={},
>               **kwargs):
>     if isinstance(func, Partial):
>       raise TypeError
>     retval = super().__new__(cls, func, *args, **kwargs)
>     retval.callable_is_static = callable_is_static
>     retval.static_argnums = set(static_argnums)
>     retval.static_kwargs = static_kwargs
>     return retval
>
>   def tree_flatten(self):
>     static_args = []
>     tree_args = []
>
>     def _append(is_static, value):
>       if is_static:
>         static_args.append(value)
>       else:
>         tree_args.append(value)
>
>     _append(self.callable_is_static, self.func)
>     for i, value in enumerate(self.args):
>       _append(i in self.static_argnums, value)
>
>     return ((list(reversed(tree_args)), self.keywords),
>             (self.callable_is_static, self.static_argnums,
>              list(reversed(static_args)), self.static_kwargs))
>
>   @classmethod
>   def tree_unflatten(cls, static, trees):
>     callable_is_static, static_argnums, static_args, static_kwargs = static
>     tree_args, tree_kwargs = trees
>     args = []
>
>     for i in range(len(static_args) + len(tree_args)):
>       if i == 0:
>         is_static = callable_is_static
>       else:
>         is_static = i - 1 in static_argnums
>       if is_static:
>         args.append(static_args.pop(0))
>       else:
>         args.append(tree_args.pop(0))
>
>     return Partial(*reversed(args),
>                    callable_is_static=callable_is_static,
>                    static_argnums=static_argnums,
>                    static_kwargs=static_kwargs,
>                    **tree_kwargs)
>
> Still getting leaked tracers, but it's getting farther than before.
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/google/jax/issues/3667#issuecomment-656240407>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAJJFVUD3FBFP2JTKTTS7ODR2XZHNANCNFSM4ORJUZEQ>
> .
>


Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -45,7 +45,7 @@
                        donation_vector, rebase_donate_argnums)
 from .tree_util import (tree_map, tree_flatten, tree_unflatten, tree_structure,
                         tree_transpose, tree_leaves, tree_multimap,
-                        treedef_is_leaf)
+                        treedef_is_leaf, Partial)
 from .util import (unzip2, curry, partial, safe_map, safe_zip, prod,
                    split_list, extend_name_stack, wrap_name)
 from .lib import xla_bridge as xb
@@ -1488,7 +1488,7 @@ def _check_inexact_input_vjp(x):
            "or complex type, got type {}")
     raise TypeError(msg.format(aval.dtype.name))
 
-def _vjp_pullback_wrapper(fun, cotangent_dtypes, io_tree, py_args):
+def _vjp_pullback_wrapper(cotangent_dtypes, io_tree, fun, py_args):
   in_tree_expected, out_tree = io_tree
   args, in_tree = tree_flatten(py_args)
   if in_tree != in_tree_expected:
@@ -1563,8 +1563,12 @@ def _vjp(fun: lu.WrappedFun, *primals, **kwargs):
     out_primal, out_vjp, aux = ad.vjp(flat_fun, primals_flat, has_aux=True)
     out_tree, aux_tree = out_aux_trees()
   out_primal_py = tree_unflatten(out_tree, out_primal)
-  vjp_py = partial(_vjp_pullback_wrapper, out_vjp,
-                   [_dtype(x) for x in out_primal], (out_tree, in_tree))
+  # Ensure that vjp_py is a PyTree so that we can pass it from the forward to the
+  # backward pass in a custom VJP.
+  vjp_py = Partial(partial(_vjp_pullback_wrapper,
+                           [_dtype(x) for x in out_primal],
+                           (out_tree, in_tree)),
+                   out_vjp)
   if not has_aux:
     return out_primal_py, vjp_py
   else:
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -27,7 +27,7 @@
 from ..tree_util import register_pytree_node
 from .. import linear_util as lu
 from ..api_util import flatten_fun, flatten_fun_nokwargs
-from ..tree_util import tree_flatten, tree_unflatten
+from ..tree_util import tree_flatten, tree_unflatten, Partial
 from .. import source_info_util
 
 zip = safe_zip
@@ -109,12 +109,16 @@ def vjp(traceable, primals, has_aux=False):
     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
   else:
     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)
-  def vjp_(*cts):
+
+  def unbound_vjp(pvals, jaxpr, consts, *cts):
     cts = tuple(map(ignore_consts, cts, pvals))
     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
     arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
     return map(instantiate_zeros, arg_cts)
 
+  # Ensure that vjp_ is a PyTree so that we can pass it from the forward to the backward
+  # pass in a custom VJP.
+  vjp_ =  Partial(partial(unbound_vjp, pvals, jaxpr), consts)
   if not has_aux:
     return out_primals, vjp_
   else:


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2863,6 +2863,38 @@ def clip_gradient(x):
 
     jax.grad(clip_gradient)(1.)  # doesn't crash
 
+  def test_nestable_vjp(self):
+    # Verify that https://github.com/google/jax/issues/3667 is resolved.
+    def f(x):
+        return x ** 2
+
+    @api.custom_vjp
+    def g(x):
+        return f(x)
+
+    def g_fwd(x):
+        y, f_vjp = api.vjp(f, x)
+        return y, f_vjp
+
+    def g_bwd(f_vjp, y_bar):
+        return f_vjp(y_bar)
+
+    g.defvjp(g_fwd, g_bwd)
+
+    # Check that VJP can be nested in simple situations.  For this to pass,
+    # vjp has to return a PyTree.
+    _, g_vjp = api.vjp(g, 1.0)
+    y, = g_vjp(1.0)
+    self.assertAllClose(y, jnp.array(2.0))
+
+    # Check that VJP can be nested in complex situations.  For this to pass,
+    # vjp can't treat the closed-over tracer x as a static argument.
+    @jit
+    def z(x):
+        _, g_vjp = api.vjp(g, x)
+        return g_vjp
+    y, = z(1.0)(3.0)
+    self.assertAllClose(y, jnp.array(6.0))
 
 class InvertibleADTest(jtu.JaxTestCase):