def use_backend(name):
  """Call fastmath functions with a specified backend."""
  if isinstance(name, Backend):
    name = name.value

  _assert_valid_backend_name(name)
  global override_backend
  prev_name_or_backend = override_backend
  override_backend = name
  # Run the decorated function in try-finally in case it throws, e.g. for tests.
  try:
    yield
  finally:
    override_backend = prev_name_or_backend
Call fastmath functions with a specified backend.
trax/fastmath/ops.py