def jit(f,
        static_argnums=(),
        xla_forced_compile=False,
        input_signature=None,
        autograph=False,
        experimental_compile=False):
  """Returns a function that runs a trace-compiled version of `f`.

  A trace-compiled version of a function `f` has the same behavior as `f` (when
  called with the same "static arguments", see below), but runs faster because
  the whole computation is compiled into a computation graph once which is
  reused for subsequent executions.

  The trace compilation happens lazily, when the returned function is called for
  the first time. The compiled function may not be cached implicitly and
  multiple calls to `jit` may not share the compiled function (see below for
  "static" vs "dynamic" arguments).

  Args:
    f: a function that takes any positional arguments `args` and any keyword
      arguments `kwargs`. `ndarray`s and things accepted by
      `tf.convert_to_tensor` in `args` and `kwargs` will be treated as 'dynamic
      arguments' in the sense that calling the function with different values
      for these arguments will not cause retracing. In contrast, arguments of
      other types in `args` and `kwargs` are treated as 'static arguments' and
      calling the function with different values of them will cause
      re-compiling. Positional arguments whose positions are in `static_argnums`
      are always treated as static arguments.
    static_argnums: a tuple of positions of arguments that will be treated as
      static arguments. Note that as aforementioned, any arguments that were not
      convertible to tensor will also be static.
    xla_forced_compile: if true, it will use XLA to force-compile the graph.
      This requires that the function only contain ops that are XLA
      compatible. It will compile the entire function into a single XLA op.
    input_signature: a list of `tf.TensorSpec`, as the input signature to
      control tracing behavior. See the
      [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
        `tf.function` for details.
    autograph: whether to use autograph to convert Python constructs such as
      `if` and `while` to their TensorFlow counterparts. See the
      [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
        `tf.function` for details.
    experimental_compile: the `experimental_compile` flag for `tf.function`. See
      the [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
      `tf.function` for details. This is the recommended way to turn on XLA for
      tf.function, but unlike xla_forced_compile, it doesn't force-compile the
      entire function into a single XLA op.

  Returns:
    A trace-compiled version of f.
  """

  @tf.function(input_signature=input_signature, autograph=autograph,
               experimental_compile=experimental_compile)
  def _tf_f(*args, **kwargs):
    """Accelerated function with tensor inputs/outputs."""
    np_args = _tf_to_np(args)
    kwargs = {k: _tf_to_np(v) for k, v in kwargs.items()}
    if xla_forced_compile:
      # Use list for mutability
      output_is_list = [False]
      output_is_empty = [False]
      output_structure = [None]
      def recorder(args, kwargs, res):
        del args, kwargs
        # Workaround b/121383831
        output_is_list[0] = isinstance(res, list)
        # If outputs are empty, xla.compile returns an `Operation`, which we
        # don't want.
        if tf.nest.flatten(res):
          output_is_empty[0] = False
          output_structure[0] = None
        else:
          output_is_empty[0] = True
          # Without deepcopy, xla.compile will change output_structure[0] to a
          # list of `Operation`.
          output_structure[0] = copy.deepcopy(res)
        return res
      f_ = _record_result_type(recorder, f)
      np_out = tf.xla.experimental.compile(lambda: f_(*np_args, **kwargs))
      # Workaround b/121383831
      if output_is_empty[0]:
        np_out = output_structure[0]
      elif (isinstance(np_out, list) and len(np_out) == 1 and
            not output_is_list[0]):
        np_out = np_out[0]
    else:
      np_out = f(*np_args, **kwargs)
    return np_out

  def _f(*args, **kwargs):
    args = [
        _canonicalize_jit_arguments(arg) if i not in static_argnums else arg
        for i, arg in enumerate(args)
    ]
    kwargs = {k: _canonicalize_jit_arguments(v) for k, v in kwargs.items()}
    tf_out = _tf_f(*args, **kwargs)
    return _tf_to_np(tf_out)

  _f.tf_function = _tf_f

  return _f
Returns a function that runs a trace-compiled version of `f`.

  A trace-compiled version of a function `f` has the same behavior as `f` (when
  called with the same "static arguments", see below), but runs faster because
  the whole computation is compiled into a computation graph once which is
  reused for subsequent executions.

  The trace compilation happens lazily, when the returned function is called for
  the first time. The compiled function may not be cached implicitly and
  multiple calls to `jit` may not share the compiled function (see below for
  "static" vs "dynamic" arguments).

  Args:
    f: a function that takes any positional arguments `args` and any keyword
      arguments `kwargs`. `ndarray`s and things accepted by
      `tf.convert_to_tensor` in `args` and `kwargs` will be treated as 'dynamic
      arguments' in the sense that calling the function with different values
      for these arguments will not cause retracing. In contrast, arguments of
      other types in `args` and `kwargs` are treated as 'static arguments' and
      calling the function with different values of them will cause
      re-compiling. Positional arguments whose positions are in `static_argnums`
      are always treated as static arguments.
    static_argnums: a tuple of positions of arguments that will be treated as
      static arguments. Note that as aforementioned, any arguments that were not
      convertible to tensor will also be static.
    xla_forced_compile: if true, it will use XLA to force-compile the graph.
      This requires that the function only contain ops that are XLA
      compatible. It will compile the entire function into a single XLA op.
    input_signature: a list of `tf.TensorSpec`, as the input signature to
      control tracing behavior. See the
      [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
        `tf.function` for details.
    autograph: whether to use autograph to convert Python constructs such as
      `if` and `while` to their TensorFlow counterparts. See the
      [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
        `tf.function` for details.
    experimental_compile: the `experimental_compile` flag for `tf.function`. See
      the [doc](https://www.tensorflow.org/api_docs/python/tf/function]) of
      `tf.function` for details. This is the recommended way to turn on XLA for
      tf.function, but unlike xla_forced_compile, it doesn't force-compile the
      entire function into a single XLA op.

  Returns:
    A trace-compiled version of f.
  
trax/tf_numpy/extensions/extensions.py