def backend(name='jax'):
  """Returns the backend used to provide fastmath ops ('tf' or 'jax')."""
  if override_backend:
    return _get_backend_from_string(override_backend)

  if default_backend:
    return _get_backend_from_string(default_backend)

  if isinstance(name, Backend):
    return _backend_dict[name]

  # name is a string.
  return _get_backend_from_string(name)
Returns the backend used to provide fastmath ops ('tf' or 'jax').
trax/fastmath/ops.py