code: 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`.
  
      `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 variable and input 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, vjp_fn = lift.vjp(partial(f, features=features), scope, x)
          return y, vjp_fn
  
        def bwd(features, vjp_fn, y_t):
          params_t, *inputs_t = vjp_fn(y_t)
          params_t = jax.tree_util.tree_map(jnp.sign, params_t)
          return (params_t, *inputs_t)
  
        dense_sign_grad = lift.custom_vjp(
            f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,))
  
      Args:
        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 passed to
          `backward_fn`.
        backward_fn: arguments are passed as (*nondiff_args, residuals, tangents)
          The function should return a tuple containing the tangents for the
          variable in the collections specified by `grad_vars` and the input
          arguments (except the scope and nondiff args).
        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 inner(
          scope_fn, repack_fn, variable_groups, rng_groups, *args
      ):
          grad_variables, other_variables = variable_groups
          scopes_treedef = None
  
          def f(grad_variables, *args):
              scope = scope_fn(
                  (grad_variables, other_variables), rng_groups
              )
              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):
              nonlocal scopes_treedef
              scopes = scope_fn(
                  (grad_variables, other_variables), rng_groups
              )
              scopes_treedef = jax.tree_util.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:
              ]  # pylint: disable=unbalanced-tuple-unpacking
              g_y, _ = g
              var_t, *inputs_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, *inputs_t)
  
          f.defvjp(f_fwd, f_bwd)
  
          return f(grad_variables, *args)
  
      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,
          name="custom_vjp",
      )
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 19 for calling a global function with example usage: custom_vjp(fn, forward_fn, backward_fn, grad_vars, nondiff_argnums) and returns: pack, f, y,  vars_out, y,  vars_out,  res, var_t,  unknown