Repository: google/flax

Issue:
Implement custom vjp
1. refactor lift.custom_vjp so the backward pass is well-defined
2. add custom_vjp to linen transforms

Fixes #1738 



Discussion:

Resolving patch:
diff --git a/flax/core/lift.py b/flax/core/lift.py
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -719,70 +719,100 @@ def scanned(broadcast_vars, carry, scan_variable_groups, rng_groups, args):
       name='scan')
 
 
-def custom_vjp(fn: Callable[..., Any], backward_fn: Callable[..., Any],
-               grad_kind: CollectionFilter = 'params',
+def custom_vjp(fn: Callable[..., Any], 
+               forward_fn: Callable[..., Any],
+               backward_fn: Callable[..., Any],
+               grad_vars: CollectionFilter = 'params',
                nondiff_argnums=()):
-  """"Lifted version of `jax.custom_vjp`.
+  """Lifted version of `jax.custom_vjp`.
 
-  `backward_fn` defines a custom vjp (backward gradient) for `fn`.
+  `forward_fn` and `backward_fn` together define a custom vjp for `fn`.
+  The original `fn` will run in case a vjp (backward gradient) is not computed.
+
+  The `forward_fn` receives the same arguments as `fn` but is expected to return
+  a tuple containing the output of `fn(scope, *args)` and the residuals that are
+  passed to `backward_fn`.
+
+  The `backward_fn` receives the nondiff arguments, residuals, and the output tangents.
+  It should return a tuple containing the input and variable tangents.
+
+  Note that the vjp function returned by `lift.vjp` can be passed as residual and
+  used in the `backward_fn`. The scope is unavailable during the backward pass.
+  If the scope is required in `backward_fn`, a snapshot of the variables can be
+  taken and returned as a residual in the `forward_fn`.
 
   Example::
 
+    f = nn.dense
+
     def fwd(scope, x, features):
-      y = nn.dense(scope, x, features)
-      return y, x
+      y, vjp_fn = lift.vjp(partial(f, features=features), scope, x)
+      return y, vjp_fn
 
-    def bwd(features, scope_fn, params, res, g):
-      x = res
-      fn = lambda params, x: nn.dense(scope_fn(params), x, features)
-      _, pullback = jax.vjp(fn, params, x)
-      g_param, g_x = pullback(g)
-      g_param = jax.tree_map(jnp.sign, g_param)
-      return g_param, g_x
+    def bwd(features, vjp_fn, y_t):
+      input_t, params_t = vjp_fn(y_t)
+      params_t = jax.tree_map(jnp.sign, params_t)
+      return input_t, params_t
 
-    dense_sign_grad = lift.custom_vjp(fwd, backward_fn=bwd, nondiff_argnums=(2,))
+    dense_sign_grad = lift.custom_vjp(
+        f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,))
 
   Args:
-    fn: should return a tuple of output and auxiliary data for the backward pass.
-    backward_fn: arguments are passed as (*nondiff_args, scope_fn, grad_variables, aux, g_y)
-      where scope_fn takes grad_variables to create the scope,
-      aux is the auxiliary data returned by `fn`,
-      and g_y is the tangent of y.
+    fn: The function to define a custom_vjp for. The first argument
+      should be a ``Module`` instance.
+    forward_fn: A function with the same arguments as `fn` returning an tuple
+      with the original output and the residuals that will be passsed to
+      `backward_fn`.
+    backward_fn: arguments are passed as (*nondiff_args, residuals, tangents)
+      The function should return a tuple containing the tangents for the
+      input arguments (except the scope and nondiff args) and the variable
+      tangents for the collections specified by `grad_vars`.
+    grad_vars: The collections for which a vjp will be computed
+      (default: "params").
+    nondiff_argnums: arguments for which no vjp is computed.
+  Returns:
+    A function with the same signature as `fn` with the custom vjp.
   """
-  # TODO(jheek) is this transform general/flexible enough?
   def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
     grad_variables, other_variables = variable_groups
-
-    def simple_scope_fn(grad_variables):
-      grad_variables = tuple(freeze(x) for x in grad_variables)
-      return scope_fn((grad_variables, other_variables), rng_groups)
+    scopes_treedef = None
 
     def f(grad_variables, *args):
       scope = scope_fn((grad_variables, other_variables), rng_groups)
-      y, _ = fn(scope, *args)
+      y = fn(scope, *args)
       vars_out = repack_fn(scope)
       return y, vars_out
     f = jax.custom_vjp(f, nondiff_argnums=nondiff_argnums)
 
     def f_fwd(grad_variables, *args):
-      scope = simple_scope_fn(grad_variables)
-      y, res = fn(scope, *args)
-      vars_out = repack_fn(scope)
-      return (y, vars_out), (res, grad_variables)
+      nonlocal scopes_treedef
+      scopes = scope_fn((grad_variables, other_variables), rng_groups)
+      scopes_treedef = jax.tree_structure(scopes)
+      y, res = forward_fn(scopes, *args)
+      vars_out = repack_fn(scopes)
+      return (y, vars_out), res
 
     def f_bwd(*args):
+      # the backward function does not pass a lifted scope
+      # to the user. Currently, there is no way to have
+      # side effects flow out of backward pass.
+      # Even without mutation variables would be ill-defined.
+      # For example, would we take a snapshot of the variables
+      # before or after calling `forward_fn`?
       nondiff_args = args[:-2]
       res, g = args[-2:]
       g_y, _ = g
-      user_res, grad_variables = res
-      return backward_fn(*nondiff_args, simple_scope_fn, grad_variables, user_res, g_y)
+      input_t, var_t = backward_fn(*nondiff_args, res, g_y)
+      assert scopes_treedef is not None, 'backward called before forward?!'
+      var_t = tuple(scopes_treedef.flatten_up_to(var_t))
+      return var_t, input_t
 
     f.defvjp(f_fwd, f_bwd)
 
     return f(grad_variables, *args)
 
-  variable_in_groups = (grad_kind, True,)
-  variable_out_groups = (grad_kind, True,)
+  variable_in_groups = (grad_vars, True)
+  variable_out_groups = (grad_vars, True)
   rng_groups = (True,)
   return pack(
       inner, variable_in_groups, variable_out_groups, rng_groups,
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -32,7 +32,7 @@
 from .pooling import avg_pool, max_pool
 from .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell
 from .stochastic import Dropout
-from .transforms import jit, named_call, checkpoint, remat, remat_scan, scan, vmap, map_variables, vjp, jvp
+from .transforms import jit, named_call, checkpoint, remat, remat_scan, scan, vmap, map_variables, vjp, jvp, custom_vjp
 from .initializers import zeros, ones
 
 # pylint: enable=g-multiple-import
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -865,6 +865,93 @@ def f(scope, x):
       rngs=rngs)
 
 
+# a version of lift.custom_vjp with a single scope function
+# this avoids having to lift multiple functions in
+# lift_transform.
+def _custom_vjp_single_scope_fn(
+    fn: Callable[..., Any],
+    backward_fn: Callable[..., Any],
+    grad_vars: lift.CollectionFilter = 'params',
+    nondiff_argnums=()):
+  nodiff_fn = functools.partial(fn, needs_residual=False)
+  forward_fn = functools.partial(fn, needs_residual=True)
+  return lift.custom_vjp(
+    nodiff_fn, forward_fn, backward_fn,
+    grad_vars, nondiff_argnums)
+
+
+def custom_vjp(fn: Callable[..., Any],
+    forward_fn: Callable[..., Any],
+    backward_fn: Callable[..., Any],
+    grad_vars: lift.CollectionFilter = 'params',
+    nondiff_argnums=()):
+  """Lifted version of `jax.custom_vjp`.
+
+  `forward_fn` and `backward_fn` together define a custom vjp for `fn`.
+  The original `fn` will run in case a vjp (backward gradient) is not computed.
+
+  The `forward_fn` receives the same arguments as `fn` but is expected to return
+  a tuple containing the output of `fn(mdl, *args)` and the residuals that are
+  passed to `backward_fn`.
+
+  The `backward_fn` receives the nondiff arguments, residuals, and the output
+  tangents. It should return a tuple containing the input and variable tangents.
+
+  Note that the vjp function returned by `nn.vjp` can be passed as residual and
+  used in the `backward_fn`. The scope is unavailable during the backward pass.
+  If the module is required in `backward_fn`, a snapshot of the variables can
+  be taken and returned as a residual in the `forward_fn`.
+
+  Example::
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        def f(mdl, x):
+          return mdl(x)
+
+        def fwd(mdl, x):
+          return nn.vjp(f, mdl, x)
+
+        def bwd(vjp_fn, y_t):
+          input_t, params_t = vjp_fn(y_t)
+          params_t = jax.tree_map(jnp.sign, params_t)
+          return input_t, params_t
+
+        sign_grad = nn.custom_vjp(
+            f, forward_fn=fwd, backward_fn=bwd)
+        return sign_grad(nn.Dense(1), x).reshape(())
+
+    x = jnp.ones((2,))
+    variables = Foo().init(random.PRNGKey(0), x)
+    grad = jax.grad(Foo().apply)(variables, x)
+
+  Args:
+    fn: The function to define a custom_vjp for.
+    forward_fn: A function with the same arguments as `fn` returning an tuple
+      with the original output and the residuals that will be passsed to
+      `backward_fn`.
+    backward_fn: arguments are passed as (*nondiff_args, residuals, tangents)
+      The function should return a tuple containing the tangents for the
+      input arguments (except the module and nondiff args) and the variable
+      tangents for the collections specified by `grad_vars`.
+    grad_vars: The collections for which a vjp will be computed
+      (default: "params").
+    nondiff_argnums: arguments for which no vjp is computed.
+  Returns:
+    A function with the same signature as `fn` with the custom vjp.
+  """
+  def shared_forward_fn(*args, needs_residual, **kwargs):
+    if needs_residual:
+      return forward_fn(*args, **kwargs)
+    else:
+      return fn(*args, ** kwargs)
+  return decorator_lift_transform(
+      _custom_vjp_single_scope_fn, shared_forward_fn,
+      backward_fn=backward_fn, grad_vars=grad_vars,
+      nondiff_argnums=nondiff_argnums,
+      multi_scope=False)
+
 
 # Special case of decorator_lift_transform to handle named calls for profiling.
 def named_call(class_fn, force=True):


Test patch:
diff --git a/tests/core/design/core_custom_vjp_test.py b/tests/core/design/core_custom_vjp_test.py
--- a/tests/core/design/core_custom_vjp_test.py
+++ b/tests/core/design/core_custom_vjp_test.py
@@ -14,6 +14,7 @@
 
 
 from typing import Sequence, Callable
+from functools import partial
 
 from absl.testing import absltest
 
@@ -29,20 +30,21 @@ def mlp_custom_grad(scope: Scope, x: Array,
                     sizes: Sequence[int] = (8, 1),
                     act_fn: Callable[[Array], Array] = nn.relu):
 
+  f = nn.dense
+
   def fwd(scope, x, features):
-    y = nn.dense(scope, x, features)
-    return y, x
+    y, vjp_fn = lift.vjp(partial(f, features=features), scope, x)
+    return y, vjp_fn
 
-  def bwd(features, scope_fn, params, res, g):
-    x = res
-    fn = lambda params, x: nn.dense(scope_fn(params), x, features)
-    _, pullback = jax.vjp(fn, params, x)
-    g_param, g_x = pullback(g)
-    g_param = jax.tree_map(jnp.sign, g_param)
-    return g_param, g_x
+  def bwd(features, res, y_t):
+    del features
+    vjp_fn = res
+    input_t, params_t = vjp_fn(y_t)
+    params_t = jax.tree_map(jnp.sign, params_t)
+    return input_t, params_t
 
   dense_custom_grad = lift.custom_vjp(
-      fwd, backward_fn=bwd, nondiff_argnums=(2,))
+      f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,))
 
   # hidden layers
   for size in sizes[:-1]:
diff --git a/tests/linen/linen_transforms_test.py b/tests/linen/linen_transforms_test.py
--- a/tests/linen/linen_transforms_test.py
+++ b/tests/linen/linen_transforms_test.py
@@ -1124,6 +1124,33 @@ def __call__(self, x):
       np.testing.assert_array_equal(vs_new['muts']['b']['outer_c']['v'],
                                     jnp.array([1.], jnp.float32))
 
+  def test_custom_vjp(self):
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        def f(mdl, x):
+          return mdl(x)
+
+        def fwd(mdl, x):
+          return nn.vjp(f, mdl, x)
+
+        def bwd(vjp_fn, y_t):
+          input_t, params_t = vjp_fn(y_t)
+          params_t = jax.tree_map(jnp.sign, params_t)
+          return input_t, params_t
+
+        sign_grad = nn.custom_vjp(
+            f, forward_fn=fwd, backward_fn=bwd)
+        return sign_grad(nn.Dense(1), x).reshape(())
+    x = jnp.ones((2,))
+    variables = Foo().init(random.PRNGKey(0), x)
+    grad = jax.grad(Foo().apply)(variables, x)
+    for grad_leaf in jax.tree_leaves(grad):
+      self.assertTrue(jnp.all(jnp.abs(grad_leaf) == 1.))
+      
+
+    
 
 if __name__ == '__main__':
   absltest.main()