Repository: google/jax

Issue:
Hiding constants from closure_convert and the AD system
Use case: Higher order functions (like ODE solvers and optimizers) with custom VJP rules.

Consider this example form the JAX documentation for `closure_convert`:

```
def minimize(objective_fn, x0):
  converted_fn, aux_args = closure_convert(objective_fn, x0)
  return _minimize(converted_fn, x0, *aux_args)

@partial(custom_vjp, nondiff_argnums=(0,))
def _minimize(objective_fn, x0, *args):
  z = objective_fn(x0, *args)
  # ... find minimizer x_opt ...
  return x_opt

def fwd(objective_fn, x0, *args):
  y = _minimize(objective_fn, x0, *args)
  return y, (y, args)

def rev(objective_fn, res, g):
  y, args = res
  y_bar = g
  # ... custom reverse-mode AD ...
  return x0_bar, *args_bars

_minimize.defvjp(fwd, rev)
```

Imagine now that `objective_fn` depends on a large matrix of constants `c`,  e.g. captured via closure. Then `closure_convert` pulls this matrix out and makes it an explicit argument of `converted_fn`.
During the reverse pass, `c_bar` has to be calculated as part of `args_bar`. This computation can be expensive and complex but it is not needed if differentiation with respect to `c` is never required or does not make sense.
There should be a way for the author of `objective_fn` to annotate `c` accordingly. This cannot be done at the level of the minimizer, where it is not known which `args` are needed in the VJP and which ones are not.

In simple cases, it is conceivable that XLA manages to optimize out the unused `c_bar`. That does not work for more complex cases, e.g. [the backward pass of an ODE integrator that has to integrate an augmented system of ODEs](https://github.com/google/jax/blob/3ac809ede30d961b873ca62f9828090709ced6be/jax/experimental/ode.py#L221).

In our concrete application (ODE integration where the dynamics function captures a large matrix), the impact of calculating `c_bar` is such that XLA consumes >44 GB of RAM during the compilation stage while attempting to generate code for the Jacobian of the dynamics function with respect to each and every element of `c`.

Can `closure_convert` be modified to not extract certain captured variables? At the moment, it ignores integer arrays, for example.
I also see parallels to https://github.com/google/jax/issues/5913, which is more focused on the JVP case.


Discussion:
The `custom_vjp` convention is to return `None` for a symbolic zero cotangent. See [the example here](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html#gradient-clipping). This would amount to setting the `c_bar` that you describe to `None` instead of filling a large matrix.

In principle, this would still define a cotangent (at very low cost), even in cases where there ought to be no meaningful derivative at all (zero or otherwise). Still, would this suffice as a workaround in your case?

Even if so, possible improvements from our end:
1. Mention the `None` convention in the `custom_vjp` docstring directly.
2. Allow for custom derivative arguments that do not participate in differentiation.

The custom VJP is defined on `_minimize`. The author of `_minimize` or some other higher-order function in a library cannot determine for which elements in `args` to return `None`. As `closure_convert` has to be used for correctness, the forward pass is handed a fairly opaque list of `args`.

Only the author of `objective_fn`, i.e. the user of the library, can determine which arguments or captured values should be ignored.
For comparison: [PyTorch autograd allows tensors to be excluded from the autograd DAG](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html#exclusion-from-the-dag).

Arguably, JAX autograd automatically opts in every floating point tensor. For the scenario described here, I see two solutions:

- An opt-out facility for the matrix of constants `c`.
- I believe that tracing should reveal that `c_bar` is unused. If that information could be made available to the backward pass, `None` could be returned for the appropriate arguments.
I wonder if it would help to dynamically create the `custom_vjp` decorated function _inside_ `minimize`? This might even be reasonably maintainable, as long as the library provides suitable helper functions for implementing the gradient rules.
> For comparison: [PyTorch autograd allows tensors to be excluded from the autograd DAG](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html#exclusion-from-the-dag).
> 
> Arguably, JAX autograd automatically opts in every floating point tensor. For the scenario described here, I see two solutions:
> 
> * An opt-out facility for the matrix of constants `c`.
> * I believe that tracing should reveal that `c_bar` is unused. If that information could be made available to the backward pass, `None` could be returned for the appropriate arguments.

JAX does have functionality like this, via `lax.stop_gradient`.

The problem is that it only works for normal JAX autodiff rules, i.e., with VJP implemented via transpose of JVP. This currently requires the lower-level JAX primitive API.
If possible, let's separate:
1. concerns about AD efficiency (unwittingly computing a cotangent that is never read downstream); from
2. concerns about AD correctness (a cotangent is mathematically ill-defined, but jax still requires a definition); from
3. limitations of the `closure_convert` utility.

Is this request captured entirely by any one of the above? A complete example would help.

Based on more recent comments, my new guess is that this is scoped to item 3. Today's `closure_convert` extracts all floating-point closure-captured arrays as a heuristic. If `closure_convert` hoisted only arguments that could possibly participate in differentiation instead, this would reduce the concerns of items 1 and 2 when using `closure_convert` together with `custom_vjp`.  Does that sound right?

If so, then that's what [this TODO comment](https://github.com/google/jax/blob/ad342419b87fbe08afe9a5f959704dee3a654a0d/jax/custom_derivatives.py#L825-L827) from @mattjj and me is about, and we can use this issue to track it.

From my point of view:
The need to use `closure_convert` and its limitations (3) are negatively affecting AD efficiency (1).
OK, let's track that here then. I suspect it will be easiest to tackle `closure_convert`'s heuristic first. If not, we can consider our options for changing the calling convention of a `custom_jvp`'s forward and reverse steps, or maybe improving DCE in AD somehow.

By the way, #5913 is about the JVP analogue of (2).
Have any paths forward to modify `closure_convert` been identified?

Resolving patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -19,7 +19,6 @@
 from typing import Callable, Generic, Optional, Sequence, Tuple, TypeVar, Any
 
 from jax import core
-from jax._src import dtypes
 from jax import linear_util as lu
 from jax.tree_util import (tree_flatten, tree_unflatten, tree_map,
                         tree_multimap, treedef_is_leaf, treedef_tuple,
@@ -831,18 +830,21 @@ def rev(objective_fn, res, g):
   else:
     return _closure_convert_for_avals(fun, in_tree, in_avals)
 
+def _is_perturbed(x: Any) -> bool:
+  if isinstance(x, ad.JVPTracer):
+    return True
+  elif isinstance(x, core.Tracer):
+    return any(_is_perturbed(attr) for name, attr in x._contents())
+  else:
+    return False
+
 @cache()
 def _closure_convert_for_avals(fun, in_tree, in_avals):
   wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
   jaxpr, out_pvals, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
   out_tree = out_tree()
 
-  # We only want to closure convert for constants with respect to which we're
-  # differentiating. As a proxy for that, we hoist consts with float dtype.
-  # TODO(frostig,mattjj): revise this approach
-  from jax.numpy import inexact
-  is_float = lambda c: dtypes.issubdtype(dtypes.dtype(c), inexact)
-  (closure_consts, hoisted_consts), merge = partition_list(is_float, consts)
+  (closure_consts, hoisted_consts), merge = partition_list(_is_perturbed, consts)
   num_consts = len(hoisted_consts)
 
   def converted_fun(*args_hconsts):
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -568,7 +568,7 @@ def __getattr__(self, name):
 
   def __repr__(self):
     base = pp('Traced<{}>with<{}>'.format(self.aval, self._trace))
-    contents = self._contents()
+    contents = [(name, pp(repr(attr))) for name, attr in self._contents()]
     if contents:
       base += pp('  with ') >> vcat(pp('{} = '.format(name)) >> pp_payload
                                     for name, pp_payload in contents)
@@ -576,7 +576,7 @@ def __repr__(self):
 
   def _contents(self):
     try:
-      return [(name, pp(repr(getattr(self, name)))) for name in self.__slots__]
+      return [(name, getattr(self, name)) for name in self.__slots__]
     except AttributeError:
       return ()
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -5211,40 +5211,82 @@ def f(x):
     self.assertAllClose(api.grad(api.grad(f))(3.), 1., check_dtypes=False)
 
   def test_closure_convert(self):
-    def minimize(objective_fn, x0):
-      converted_fn, aux_args = api.closure_convert(objective_fn, x0)
-      return _minimize(converted_fn, x0, *aux_args)
+    def cos_after(fn, x):
+      converted_fn, aux_args = api.closure_convert(fn, x)
+      self.assertLessEqual(len(aux_args), 1)
+      return _cos_after(converted_fn, x, *aux_args)
 
     @partial(api.custom_vjp, nondiff_argnums=(0,))
-    def _minimize(objective_fn, x0, *args):
-      _ = objective_fn(x0, *args)
-      return jnp.cos(x0)
+    def _cos_after(fn, x, *args):
+      return jnp.cos(fn(x, *args))
 
-    def fwd(objective_fn, x0, *args):
-      y = _minimize(objective_fn, x0, *args)
-      return y, (y, args)
+    def fwd(fn, x, *args):
+      y = _cos_after(fn, x, *args)
+      return y, (x, args)
 
-    def rev(objective_fn, res, g):
-      y, args = res
-      x0_bar = 17. * y
+    def rev(fn, res, g):
+      x, args = res
+      x_bar = 17. * x
       args_bars = [42. * a for a in args]
-      return (x0_bar, *args_bars)
+      return (x_bar, *args_bars)
 
-    _minimize.defvjp(fwd, rev)
+    _cos_after.defvjp(fwd, rev)
 
-    def obj(c, x):
+    def dist(c, x):
       return jnp.sum((x - c) ** 2.)
 
     def solve(c, x):
       def closure(x):
-        return obj(c, x)
-      return jnp.sum(minimize(closure, x))
+        return dist(c, x)
+      return cos_after(closure, x)
 
-    c, x = jnp.ones(2), jnp.zeros(2)
-    self.assertAllClose(solve(c, x), 2.0, check_dtypes=False)
+    c, x = 2. * jnp.ones(2), jnp.ones(2)
+    expected = jnp.cos(dist(c, x))
+    self.assertAllClose(solve(c, x), expected, check_dtypes=False)
     g_c, g_x = api.grad(solve, argnums=(0, 1))(c, x)
-    self.assertAllClose(g_c, 42. * jnp.ones(2), check_dtypes=False)
-    self.assertAllClose(g_x, 17. * jnp.ones(2), check_dtypes=False)
+    self.assertAllClose(g_c, 42. * c, check_dtypes=False)
+    self.assertAllClose(g_x, 17. * x, check_dtypes=False)
+
+  def test_closure_convert_mixed_consts(self):
+    # Like test_closure_convert, but close over values that
+    # participate in AD as well as values that do not.
+    # See https://github.com/google/jax/issues/6415
+
+    def cos_after(fn, x):
+      converted_fn, aux_args = api.closure_convert(fn, x)
+      self.assertLessEqual(len(aux_args), 1)
+      return _cos_after(converted_fn, x, *aux_args)
+
+    @partial(api.custom_vjp, nondiff_argnums=(0,))
+    def _cos_after(fn, x, *args):
+      return jnp.cos(fn(x, *args))
+
+    def fwd(fn, x, *args):
+      y = _cos_after(fn, x, *args)
+      return y, (x, args)
+
+    def rev(fn, res, g):
+      x, args = res
+      x_bar = 17. * x
+      args_bars = [42. * a for a in args]
+      return (x_bar, *args_bars)
+
+    _cos_after.defvjp(fwd, rev)
+
+    def dist(c, s, x):
+      return jnp.sum(s * (x - c) ** 2.)
+
+    def solve(c, s, x):
+      def closure(x):
+        return dist(c, s, x)
+      return cos_after(closure, x)
+
+    c, s, x = 2. * jnp.ones(2), 3. * jnp.ones(2), jnp.ones(2)
+    expected = jnp.cos(dist(c, s, x))
+    self.assertAllClose(solve(c, s, x), expected, check_dtypes=False)
+    g_c, g_x = api.grad(solve, argnums=(0, 2))(c, s, x)
+    self.assertAllClose(g_c, 42. * c, check_dtypes=False)
+    self.assertAllClose(g_x, 17. * x, check_dtypes=False)
 
 
 class CustomTransposeTest(jtu.JaxTestCase):