Repository: google/jax

Issue:
utility for closure-conversion in higher-order functions with custom derivatives
When a higher-order python function is associated with a custom derivative (e.g. via `custom_vjp`), differentiation doesn't handle the closures of its function arguments. A workaround is to accept auxiliary arguments and thread them back into input functions that otherwise would have closed over them. An example sketch:

```python
@custom_vjp
def minimize(objective_fn, x0, objective_aux_args=()):
  # ...
  t = objective_fn(..., *objective_aux_args)
  # ...
```

An alternative approach taken by `jax.experimental.ode.odeint` is to stage function arguments out to jaxprs, converting closures in the process. We could try and factor [this setup](https://github.com/google/jax/blob/943c7794f92f32c4796cf980eedd5a11c86eea66/jax/experimental/ode.py#L50) into a common utility for use in other similar applications.


Discussion:
The workaround based on auxiliary arguments has several disadvantages. It can be tricky to enforce, and it can require inconvenient code restructuring in callers. For more on why `odeint` takes the approach that it does, see #2718, #3557, #3558, and the PR resolving them: #3562
I'm currently struggling with this problem for a `custom_jvp` for a higher-order derivative from a solver.  It would be really cool to have a docstring for `closure_convert` to make it clear how to use it in other places. :-)
Since this is an issue specific to custom derivatives, I'm wondering if we ought to offer it as part of `custom_vjp`, `custom_gradient`, etc., in a way that lets you write:

```python
@partial(custom_vjp, closure_convert_argnums=0)
def minimize(objective_fn, x0):
  # ...
```

@rpadams Would that work for your solver?
@mattjj Having written `ode.closure_convert` and `custom_vjp`, what do you think?
Thanks for the fast reply @froystig! If I'm understanding the situation, I think this would resolve my situation in which the thing I'm writing a `custom_vjp` for is a solver that consumes multiple levels of closures.  Threading the arguments through this seems daunting. 

More context: I have an energy function that depends on a Jacobian (1st deriv), that is part of a Lagrangian which gives an ODE via Euler-Lagrange (2nd deriv), which I then use an implicit solver for (3rd deriv).  Then I'd like to get gradients back through that stack (4th deriv) without having to backprop through my implicit (in the "backward Euler" sense) solver using the implicit (in the "implicit function theorem" sense) gradient.  There are many closures/partials/lambdas along the way that I think are not playing nicely with the outer-loop `custom_jvp`... (However, that last level thing is not mathematically painful because I did most of the IFT work with the second-order solver.)
> Threading the arguments through this seems daunting.

Based on your context, it seems that what you'd find daunting is having to pass around extra arguments explicitly on the way down to calling the solver, to avoid forming closures anywhere. Is that correct? By contrast, would you be OK threading the arguments around _within_ the solver implementation, supposing they were extracted from the incoming closures for you?

An amendment to my previous comment: it might make more sense to return to the original idea and offer a `closure_convert` utility directly—rather than an option to `custom_vjp`—so that you can control the placement and threading of arguments in your solver implementation. In my sketch above, `minimize` has implicit arguments (via the hypothetical closure conversion) that aren't apparent in its signature.

Altogether here's roughly how using this would look:

```python
def minimize(objective_fn, x0):
  converted_fn, consts = closure_convert(objective_fn, x0)
  return _minimize(converted_fn, x0, consts)

@partial(custom_vjp, nondiff_argnums=0)
def _minimize(objective_fn, x0, objective_aux_args):
  z = objective_fn(x0, objective_aux_args)
  # ...
```

Assuming we're thinking about this the same way, it's that there are many closures going into forming the objective function minimized by the solver.  Threading arguments around just the solver itself isn't too bad --- I'm basically doing that already in order to avoid re-jitting my Levenberg-Marquardt implementation every time step.  

I should say that part of what I'm confused about is the statement in the docs:

> A limitation to this approach is that the argument f can’t close over any values involved in differentiation.

It seems like "values involved in differentiation" necessarily covers a lot of ground, i.e., essentially everything that's gone into the objective function, no?  I'm interpreting this as "you can't use closures/lambda/partials in building your objective function", but maybe that's overly broad? In my case I'm making pretty extensive use of a generator pattern, e.g., `generate_lagrangian` that hands back a Lagrangian function that I hand to a `generate_euler_lagrange` that gives a function I can hand to a `generate_time_stepper` that constructs an objective for the implicit Euler, etc.  I think you guys can appreciate that this is a pretty "SICM" kind of thing I'm doing. :-)


Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -68,7 +68,8 @@
 from .interpreters import masking
 from .interpreters import invertible_ad as iad
 from .interpreters.invertible_ad import custom_ivjp
-from .custom_derivatives import custom_jvp, custom_vjp, custom_gradient
+from .custom_derivatives import (closure_convert, custom_gradient, custom_jvp,
+                                 custom_vjp)
 from .config import flags, config, bool_env
 
 traceback_util.register_exclusion(__file__)
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -19,10 +19,11 @@
 from typing import Callable, Sequence, Tuple, Any
 
 from . import core
+from . import dtypes
 from . import linear_util as lu
 from .tree_util import (tree_flatten, tree_unflatten, tree_map, tree_multimap,
                         register_pytree_node_class)
-from .util import safe_zip, safe_map, split_list
+from .util import cache, safe_zip, safe_map, split_list
 from .api_util import flatten_fun_nokwargs, argnums_partial, wrap_hashably
 from .core import raise_to_shaped
 from .ad_util import Zero, zeros_like_aval, stop_gradient_p
@@ -837,3 +838,108 @@ def tree_flatten(self):
   def tree_unflatten(cls, aux, consts):
     jaxpr, in_tree, out_tree = aux
     return cls(jaxpr, in_tree, out_tree, consts)
+
+
+def closure_convert(fun, *example_args):
+  """Closure conversion utility, for use with higher-order custom derivatives.
+
+  To define custom derivatives such as with ``jax.custom_vjp(f)``, the target
+  function ``f`` must take, as formal arguments, all values involved in
+  differentiation. If ``f`` is a higher-order function, in that it accepts as an
+  argument a Python function ``g``, then values stored away in ``g``'s closure
+  will not be visible to the custom derivative rules, and attempts at AD
+  involving these values will fail. One way around this is to convert the
+  closure by extracting these values, and to pass them as explicit formal
+  arguments across the custom derivative boundary. This utility carries out that
+  conversion. More precisely, it closure-converts the function ``fun``
+  specialized to the types of the arguments given in ``example_args``.
+
+  When we refer here to "values in the closure" of ``fun``, we do not mean the
+  values that are captured by Python directly when ``fun`` is defined (e.g. the
+  Python objects in ``fun.__closure__``, if the attribute exists). Rather, we
+  mean values encountered during the execution of ``fun`` on ``example_args``
+  that determine its output. This may include, for instance, arrays captured
+  transitively in Python closures, i.e. in the Python closure of functions
+  called by ``fun``, the closures of the functions that they call, and so forth.
+
+  The function ``fun`` must be a pure function.
+
+  Example usage::
+
+    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)
+
+  Args:
+    fun: Python callable to be converted. Must be a pure function.
+    example_args: Arrays, scalars, or (nested) standard Python
+      containers (tuples, lists, dicts, namedtuples, i.e., pytrees)
+      thereof, used to determine the types of the formal arguments to
+      ``fun``. This type-specialized form of ``fun`` is the function
+      that will be closure converted.
+
+  """
+  flat_args, in_tree = tree_flatten(example_args)
+  in_avals = tuple(map(abstractify, flat_args))
+  return _closure_convert_for_avals(fun, in_tree, in_avals)
+
+@cache()
+def _closure_convert_for_avals(fun, in_tree, in_avals):
+  if config.omnistaging_enabled:
+    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)
+  else:
+    in_pvals = [pe.PartialVal.unknown(aval) for aval in in_avals]
+    wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
+    with core.initial_style_staging():  # type: ignore
+      jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
+        wrapped_fun, in_pvals, instantiate=True, stage_out=False)  # type: ignore
+  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(mattjj): revise this approach
+  from .numpy import inexact
+  is_float = lambda c: dtypes.issubdtype(dtypes.dtype(c), inexact)
+  (closure_consts, hoisted_consts), merge = partition_list(is_float, consts)
+  num_consts = len(hoisted_consts)
+
+  def converted_fun(*args_hconsts):
+    num_args = len(args_hconsts) - num_consts
+    args, hoisted_consts = split_list(args_hconsts, [num_args])
+    consts = merge(closure_consts, hoisted_consts)
+    all_args, in_tree2 = tree_flatten(tuple(args))
+    assert in_tree == in_tree2
+    out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
+    return tree_unflatten(out_tree, out_flat)
+
+  return converted_fun, hoisted_consts
+
+def partition_list(choice, lst):
+  out = [], []
+  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]
+  def merge(l1, l2):
+    i1, i2 = iter(l1), iter(l2)
+    return [next(i2 if snd else i1) for snd in which]
+  return out, merge
+
+def abstractify(x):
+  return core.raise_to_shaped(core.get_aval(x))
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
--- a/jax/experimental/ode.py
+++ b/jax/experimental/ode.py
@@ -32,61 +32,17 @@
 import jax
 import jax.numpy as jnp
 from jax import core
-from jax import dtypes
+from jax import custom_derivatives
 from jax import lax
-from jax.util import safe_map, safe_zip, cache, split_list
-from jax.api_util import flatten_fun_nokwargs
+from jax.util import safe_map, safe_zip
 from jax.flatten_util import ravel_pytree
-from jax.tree_util import tree_map, tree_flatten, tree_unflatten
-from jax.interpreters import partial_eval as pe
+from jax.tree_util import tree_map
 from jax import linear_util as lu
-from jax import config
 
 map = safe_map
 zip = safe_zip
 
 
-@cache()
-def closure_convert(fun, in_tree, in_avals):
-  if config.omnistaging_enabled:
-    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)
-  else:
-    in_pvals = [pe.PartialVal.unknown(aval) for aval in in_avals]
-    wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
-    with core.initial_style_staging():  # type: ignore
-      jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
-        wrapped_fun, in_pvals, instantiate=True, stage_out=False)  # type: ignore
-  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(mattjj): revise this approach
-  is_float = lambda c: dtypes.issubdtype(dtypes.dtype(c), jnp.inexact)
-  (closure_consts, hoisted_consts), merge = partition_list(is_float, consts)
-  num_consts = len(hoisted_consts)
-
-  def converted_fun(y, t, *hconsts_args):
-    hoisted_consts, args = split_list(hconsts_args, [num_consts])
-    consts = merge(closure_consts, hoisted_consts)
-    all_args, in_tree2 = tree_flatten((y, t, *args))
-    assert in_tree == in_tree2
-    out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
-    return tree_unflatten(out_tree, out_flat)
-
-  return converted_fun, hoisted_consts
-
-def partition_list(choice, lst):
-  out = [], []
-  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]
-  def merge(l1, l2):
-    i1, i2 = iter(l1), iter(l2)
-    return [next(i2 if snd else i1) for snd in which]
-  return out, merge
-
-def abstractify(x):
-  return core.raise_to_shaped(core.get_aval(x))
-
 def ravel_first_arg(f, unravel):
   return ravel_first_arg_(lu.wrap_init(f), unravel).call_wrapped
 
@@ -213,11 +169,8 @@ def _check_arg(arg):
              "\n{}.")
       raise TypeError(msg.format(arg))
 
-  flat_args, in_tree = tree_flatten((y0, t[0], *args))
-  in_avals = tuple(map(abstractify, flat_args))
-  converted, consts = closure_convert(func, in_tree, in_avals)
-
-  return _odeint_wrapper(converted, rtol, atol, mxstep, y0, t, *consts, *args)
+  converted, consts = custom_derivatives.closure_convert(func, y0, t[0], *args)
+  return _odeint_wrapper(converted, rtol, atol, mxstep, y0, t, *args, *consts)
 
 @partial(jax.jit, static_argnums=(0, 1, 2, 3))
 def _odeint_wrapper(func, rtol, atol, mxstep, y0, ts, *args):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4333,6 +4333,42 @@ def f(x):
         api.grad(lambda x: jnp.sum(jnp.sin(x)))(jnp.arange(3.)) * jnp.array([3., 4., 5.]),
         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)
+
+    @partial(api.custom_vjp, nondiff_argnums=(0,))
+    def _minimize(objective_fn, x0, *args):
+      _ = objective_fn(x0, *args)
+      return jnp.cos(x0)
+
+    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
+      x0_bar = 17. * y
+      args_bars = [42. * a for a in args]
+      return (x0_bar, *args_bars)
+
+    _minimize.defvjp(fwd, rev)
+
+    def obj(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))
+
+    c, x = jnp.ones(2), jnp.zeros(2)
+    self.assertAllClose(solve(c, x), 2.0, 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)
+
 
 class InvertibleADTest(jtu.JaxTestCase):