def jit(*args, **kwargs):
  """Just-In-Time compiles the given function for use on accelerators."""
  global _disable_jit
  if _disable_jit:
    return args[0]  # jit(f, **unused_now_jit_kwargs) = f
  return backend()['jit'](*args, **kwargs)
Just-In-Time compiles the given function for use on accelerators.
trax/fastmath/ops.py