code: def jit(
      target: Target,
      variables: lift.CollectionFilter = True,
      rngs: lift.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,
      methods=None,
  ) -> Target:
      """Lifted version of ``jax.jit``.
  
      Args:
        target: a ``Module`` or a function taking a ``Module`` as its first
          argument.
        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'``.
        methods: If ``target`` is a ``Module``, the methods of ``Module`` to jit.
  
      Returns:
        A wrapped version of target, set up for just-in-time compilation.
      """
      return lift_transform(
          lift.jit,
          target,
          variables=variables,
          rngs=rngs,
          static_argnums=static_argnums,
          static_argnames=static_argnames,
          donate_argnums=donate_argnums,
          device=device,
          backend=backend,
          methods=methods,
      )
  
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 12 for calling a global function with example usage: jit(target, variables, rngs, static_argnums, static_argnames, donate_argnums, device, backend, methods) and returns: lift_transform