code: def set_backend(backend):
      """
      Set backend for Funsor.
  
      Currently three backends are supported: "numpy" (default), "torch", and
      "jax". Funsor runs with only one backend at a time.
  
      .. note: When `jax` backend is set, we cannot revert back to the default
      `numpy` backend because we dispatch to using `jax.numpy` all ops with
      `numpy.ndarray` or `numpy.generic` inputs.
  
      :param str backend: either "numpy", "torch", or "jax".
      """
      global _FUNSOR_BACKEND, _JAX_LOADED
      if _FUNSOR_BACKEND == backend:
          return
  
      if backend == "numpy":
          if _JAX_LOADED:
              raise ValueError(
                  "Cannot revert back to NumPy backend when JAX backend has been set."
              )
          else:
              _FUNSOR_BACKEND = "numpy"
      elif backend == "torch":
          _FUNSOR_BACKEND = "torch"
  
          import torch  # noqa: F401
  
          import funsor.torch  # noqa: F401
      elif backend == "jax":
          _FUNSOR_BACKEND = "jax"
          _JAX_LOADED = True
          global _JAX_COMPILED_FUNCTION_TYPE
  
          import jax  # noqa: F401
  
          import funsor.jax  # noqa: F401
  
          if _JAX_COMPILED_FUNCTION_TYPE is None:
              _JAX_COMPILED_FUNCTION_TYPE = type(jax.jit(lambda: 0))
      else:
          raise ValueError(
              "backend should be either 'numpy', 'torch', or 'jax'"
              ", got {}".format(backend)
          )
  
imports: import functools
  import inspect
  import pprint
  import re
  import sys
  import numpy
label_desc: Use image node_id 12 for calling a global function with example usage: set_backend(backend) without return types