code: def jvp(
      fn: Callable[..., Any],
      scope: Scope,
      primals,
      tangents,
      variable_tangents,
      variables: CollectionFilter = True,
      rngs: PRNGSequenceFilter = True,
  ) -> Tuple[Any, Any]:
      """A lifted version of ``jax.jvp``.
  
      See ``jax.jvp`` for the unlifted Jacobian-vector product (forward gradient).
  
      Note that no tangents are returned for variables. When variable tangents
      are required their value should be returned explicitly by `fn`
      using `scope.variables()`.
  
      Example::
  
        def learn_scale(scope, x):
          p = scope.param('scale', nn.initializers.zeros_init(), ())
          return p * x
  
        def f(scope, x):
          vars_t = jax.tree_util.tree_map(jnp.ones_like,
                                          scope.variables().get('params', {}))
          x, out_t = lift.jvp(
              learn_scale, scope, (x,), (jnp.zeros_like(x),),
              variable_tangents={'params': vars_t})
          return out_t
  
      Args:
        fn: The function to be transformed.
        scope: The scope(s) which should be lifted into the transform.
        primals: The primal values at which the Jacobian of ``fun`` should be
          evaluated. Should be either a tuple or a list of arguments,
          and its length should be equal to the number of positional parameters of
          ``fun``.
        tangents: The tangent vector for which the Jacobian-vector product should be
          evaluated. Should be either a tuple or a list of tangents, with the same
          tree structure and array shapes as ``primals``.
        variable_tangents: A dict or PyTree fo dicts with the same structure as
          scopes. Each entry in the dict specifies the tangents for a variable
          collection. Not specifying a collection in variable_tangents is
          equivalent to passing a zero vector as the tangent.
        variables: other variables collections that are available inside `fn` but
          do not receive a tangent.
        rngs: the prngs that are available inside `fn`.
  
      Returns:
        A ``(primals_out, tangents_out)`` pair, where ``primals_out`` is
        ``fun(*primals)``, and ``tangents_out`` is the Jacobian-vector product of
        ``function`` evaluated at ``primals`` with ``tangents``. The
        ``tangents_out`` value has the same Python tree structure and shapes as
        ``primals_out``.
      """
  
      def inner(
          scope_fn, repack_fn, variable_groups, rng_groups, *args
      ):
          jvp_vars, other_vars = variable_groups
  
          @functools.wraps(fn)
          def wrapper(vars_primals, args):
              variable_groups = (vars_primals, other_vars)
              scope = scope_fn(variable_groups, rng_groups)
              y = fn(scope, *args)
              return y, repack_fn(scope)
  
          (y, out_vars), out_tangents = jax.jvp(
              wrapper, (jvp_vars, args), (variable_tangents, tangents)
          )
          return (y, out_tangents[0]), out_vars
  
      # filter out empty tangent collections because JAX will error on non-equal
      # tree structure for example: {"params": {}} != {}.
      treedef = jax.tree_util.tree_structure(scope)
  
      variable_tangents = tuple(
          {
              k: v for k, v in vt.items() if v
          }  # pylint: disable=g-complex-comprehension
          for vt in treedef.flatten_up_to(variable_tangents)
      )
      target = tuple(variable_tangents[0].keys())
      return pack(
          inner,
          (target, variables),
          (variables,),
          (rngs,),
          name="jvp",
          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 13 for calling a global function with example usage: jvp(fn, scope, primals, tangents, variable_tangents, variables, rngs) and returns: pack, y,  out_tangents,  out_vars, y,  repack_fn