code: def vjp(
      fn: Callable[..., Any],
      scope: Scope,
      *primals,
      has_aux: bool = False,
      reduce_axes=(),
      vjp_variables: CollectionFilter = "params",
      variables: CollectionFilter = True,
      rngs: PRNGSequenceFilter = True,
  ) -> Union[
      Tuple[Any, Callable[..., Any]],
      Tuple[Any, Callable[..., Any], Any],
  ]:
      """A lifted version of ``jax.vjp``.
  
      See ``jax.vjp`` for the unlifted vector-Jacobian product (backward gradient).
  
      Note that a gradient is returned for all variables in the collections
      specified by `vjp_variables`. However, the backward function only expects
      a cotangent for the return value of `fn`. If variables require a co-tangent
      as well they can be returned from `fn` using `scope.variables()`.
  
      Example::
  
        def learn_scale(scope, x, y):
          p = scope.param('scale', nn.initializers.zeros_init(), ())
          return p * x * y
        def f(scope, x, y):
          z, bwd = lift.vjp(learn_scale, scope, x, y)
          params_grad, x_grad, y_grad = bwd(jnp.ones(z.shape))
          return z, params_grad, x_grad, y_grad
  
      Args:
        fn: Function to be differentiated. Its arguments should be arrays, scalars,
          or standard Python containers of arrays or scalars. It should return an
          array, scalar, or standard Python container of arrays or scalars. It will
          receive the scope and primals as arguments.
        scope: The scope of which the variables will be differentiated.
        *primals: A sequence of primal values at which the Jacobian of ``fn``
          should be evaluated. The length of ``primals`` should be equal to the
          number of positional parameters to ``fn``. Each primal value should be a
          tuple of arrays, scalar, or standard Python containers thereof.
        has_aux: Optional, bool. Indicates whether ``fn`` returns a pair where the
         first element is considered the output of the mathematical function to be
         differentiated and the second element is auxiliary data. Default ``False``.
        reduce_axes: Optional, tuple of axis names. If an axis is listed here, and
          ``fn`` implicitly broadcasts a value over that axis, the backward pass
          will perform a ``psum`` of the corresponding gradient. Otherwise, the
          VJP will be per-example over named axes. For example, if ``'batch'``
          is a named batch axis, ``vjp(f, *args, reduce_axes=('batch',))`` will
          create a VJP function that sums over the batch while ``vjp(f, *args)``
          will create a per-example VJP.
        vjp_variables: The vjpfun will return a cotangent vector for all
          variable collections specified by this filter.
        variables: other variables collections that are available inside `fn` but
          do not receive a cotangent.
        rngs: the prngs that are available inside `fn`.
  
      Returns:
        If ``has_aux`` is ``False``, returns a ``(primals_out, vjpfun)`` pair, where
        ``primals_out`` is ``fn(*primals)``.
        ``vjpfun`` is a function from a cotangent vector with the same shape as
        ``primals_out`` to a tuple of cotangent vectors with the same shape as
        ``primals``, representing the vector-Jacobian product of ``fn`` evaluated at
        ``primals``. If ``has_aux`` is ``True``, returns a
        ``(primals_out, vjpfun, aux)`` tuple where ``aux`` is the auxiliary data
        returned by ``fn``.
      """
  
      def inner(
          scope_fn, repack_fn, variable_groups, rng_groups, *args
      ):
          vjp_vars, other_vars = variable_groups
  
          @functools.wraps(fn)
          def wrapper(vjp_vars, *args):
              variable_groups = (vjp_vars, other_vars)
              scope = scope_fn(variable_groups, rng_groups)
              if has_aux:
                  y, aux = fn(scope, *args)
              else:
                  y = fn(scope, *args)
                  aux = ()
              return y, (aux, repack_fn(scope))
  
          y, bwd, (aux, out_vars) = jax.vjp(
              wrapper,
              vjp_vars,
              *args,
              reduce_axes=reduce_axes,
              has_aux=True,
          )
          treedef = jax.tree_util.tree_structure(scope)
          bwd = jax.tree_util.Partial(
              functools.partial(_bwd_wrapper, treedef), bwd
          )
          if has_aux:
              return (y, bwd, aux), out_vars
          else:
              return (y, bwd), out_vars
  
      return pack(
          inner,
          (vjp_variables, variables),
          (variables,),
          (rngs,),
          name="vjp",
          enable_kwargs=False,
      )(scope, *primals)
  
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 11 for calling a global function with example usage: vjp(fn, scope) and returns: pack, y,  aux,  repack_fn, y,  bwd,  aux,  out_vars, y,  bwd,  out_vars