code: 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 variable and input 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::
  
        >>> import flax.linen as nn
        >>> import jax, jax.numpy as jnp
  
        >>> 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):
        ...       params_t, *inputs_t = vjp_fn(y_t)
        ...       params_t = jax.tree_util.tree_map(jnp.sign, params_t)
        ...       return (params_t, *inputs_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(jax.random.key(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 variable in the collections
          specified by ``grad_vars`` and the input arguments (except the module 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 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,
      )
  
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
  from flax import errors, struct, traceback_util
  from flax.core import Scope, lift, meta
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import module
  from flax.linen.module import Module, Variable, _derive_profiling_name, _get_unbound_fn, wrap_method_once
  import jax
label_desc: Use image node_id 27 for calling a global function with example usage: custom_vjp(fn, forward_fn, backward_fn, grad_vars, nondiff_argnums) and returns: decorator_lift_transform, forward_fn, fn