code: def jit(
      fn: Callable[..., Any],
      variables: CollectionFilter = True,
      rngs: PRNGSequenceFilter = True,
      static_argnums: Union[int, Iterable[int]] = (),
      static_argnames: Union[str, Iterable[str]] = (),
      donate_argnums: Union[int, Iterable[int]] = (),
      device=None,
      backend: Union[str, None] = None,
  ) -> Callable[..., Any]:
      """Lifted version of ``jax.jit``.
  
      Args:
        fn: Scope function to be jitted.
        variables: The variable collections that are lifted. By default all
          collections are lifted.
        rngs: The PRNG sequences that are lifted. By default all PRNG sequences
          are lifted.
        static_argnums: An int or collection of ints specifying which positional
          arguments to treat as static (compile-time constant). Operations that only
          depend on static arguments will be constant-folded in Python (during
          tracing), and so the corresponding argument values can be any Python
          object. Static arguments should be hashable, meaning both ``__hash__`` and
          ``__eq__`` are implemented, and immutable. Calling the jitted function
          with different values for these constants will trigger recompilation. If
          the jitted function is called with fewer positional arguments than
          indicated by ``static_argnums`` then an error is raised. Arguments that
          are not arrays or containers thereof must be marked as static.
          Defaults to ().
        static_argnames: An optional string or collection of strings specifying
          which named arguments to treat as static (compile-time constant). See the
          comment on ``static_argnums`` for details. If not
          provided but ``static_argnums`` is set, the default is based on calling
          ``inspect.signature(fun)`` to find corresponding named arguments.
        donate_argnums: Specify which arguments are "donated" to the computation.
          It is safe to donate arguments if you no longer need them once the
          computation has finished. In some cases XLA can make use of donated
          buffers to reduce the amount of memory needed to perform a computation,
          for example recycling one of your input buffers to store a result. You
          should not reuse buffers that you donate to a computation, JAX will raise
          an error if you try to.
        device: This is an experimental feature and the API is likely to change.
          Optional, the Device the jitted function will run on. (Available devices
          can be retrieved via :py:func:`jax.devices`.) The default is inherited
          from XLA's DeviceAssignment logic and is usually to use
          ``jax.devices()[0]``.
        backend: a string representing the XLA backend: ``'cpu'``, ``'gpu'``, or
          ``'tpu'``.
  
      Returns:
        A wrapped version of ``fn``, set up for just-in-time compilation.
      """
      if not isinstance(static_argnums, Iterable):
          static_argnums = (static_argnums,)
      if not isinstance(donate_argnums, Iterable):
          donate_argnums = (donate_argnums,)
      # offset argnums by two because first argument in the original function is the
      # scope while jitted has 3 functions before the user arguments.
      static_argnums = (0,) + tuple(
          i + 2 for i in static_argnums if i > 0
      )
      donate_argnums = tuple(i + 2 for i in donate_argnums if i > 0)
  
      # Close over scope_fn & repack_fn to avoid recompilation
      # this is impure but we use the fingerprint arg to differentiate between cases
      # where scope_fn or repack_fn actually produce non-identical results.
      scope_fn = None  # type: Optional[Callable]
      repack_fn = None  # type: Optional[Callable]
  
      @functools.partial(
          jax.jit,
          static_argnums=static_argnums,
          static_argnames=static_argnames,
          donate_argnums=donate_argnums,
          device=device,
          backend=backend,
      )
      @functools.wraps(fn)
      def jitted(
          fingerprint, variable_groups, rng_groups, *args, **kwargs
      ):
          nonlocal scope_fn, repack_fn
          # fingerprint is only used to differentiate the cache signature for cases
          # where different collections are mutable.
          del fingerprint
          scope = scope_fn(
              variable_groups, rng_groups
          )  # pylint: disable=not-callable
          y = fn(scope, *args, **kwargs)
          return y, repack_fn(scope)  # pylint: disable=not-callable
  
      def inner(
          scope_fun,
          repack_fun,
          variable_groups,
          rng_groups,
          *args,
          **kwargs,
      ):
          nonlocal scope_fn, repack_fn
          try:
              scope_fn = scope_fun
              repack_fn = repack_fun
              scopes = jax.tree_util.tree_leaves(
                  scope_fn(variable_groups, rng_groups)
              )
              mutable = tuple(
                  _hashable_filter(scope.mutable) for scope in scopes
              )
              return jitted(
                  mutable, variable_groups, rng_groups, *args, **kwargs
              )
          finally:
              scope_fn, repack_fn = None, None
  
      return pack(
          inner,
          (variables,),
          (variables,),
          (rngs,),
          name="jit",
          enable_kwargs=True,
      )
  
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 22 for calling a global function with example usage: jit(fn, variables, rngs, static_argnums, static_argnames, donate_argnums, device, backend) and returns: pack, y,  repack_fn, jitted