code: def multi_backend_test(
      globals_dict,
      relative_module_name,
      backends=("jax", "tensorflow", "numpy"),
      test_case=None,
  ):
      """Multi-backend test decorator.
  
      The end goal of this decorator is that the decorated test case is removed, and
      replaced with a set of new test cases that have been rewritten to use one or
      more backends. E.g., a test case named `Test` will by default be rewritten to
      `Test_jax`, 'Test_tensorflow' and `Test_numpy` which use the JAX, TensorFlow
      and NumPy backends, respectively.
  
      The decorator works by using the dynamic rewrite system to rewrite imports of
      the module the test is defined in, and inserting the approriately renamed test
      cases into the `globals()` dictionary of the original module. A side-effect of
      this is that the global code inside the module is run `1 + len(backends)`
      times, so avoid doing anything expensive there. This does mean that the
      original module needs to be in a runnable state, i.e., when it uses symbols
      from `backend`, those must be actually present in the literal `backend`
      module.
  
      A subtle point about what this decorator does in the rewritten modules: the
      rewrite system changes the behavior of this decorator to act as a passthrough
      to avoid infinite rewriting loops.
  
      Args:
        globals_dict: Python dictionary of strings to symbols. Set this to the value
          of `globals()`.
        relative_module_name: Python string. The module name of the module where the
          decorated test resides relative to `inference_gym`. You must not use
          `__name__` for this as that is set to a defective value of `__main__`
          which is sufficiently abnormal that the rewrite system does not work on
          it.
        backends: Python iterable of strings. Which backends to test with.
        test_case: The actual test case to decorate.
  
      Returns:
        None, to delete the original test case.
      """
      if test_case is None:
          return lambda test_case: multi_backend_test(  # pylint: disable=g-long-lambda
              globals_dict=globals_dict,
              relative_module_name=relative_module_name,
              test_case=test_case,
          )
  
      if BACKEND is not None:
          return test_case
  
      if relative_module_name == "__main__":
          raise ValueError(
              "module_name should be written out manually, not by passing __name__."
          )
  
      # This assumes `test_util` is 2 levels deep inside of `inference_gym`. If we
      # move it, we'd change the `-2` to equal the (negative) nesting level.
      root_name_comps = __name__.split(".")[:-2]
      relative_module_name_comps = relative_module_name.split(".")
  
      # Register the rewrite hooks.
      importlib.import_module(
          ".".join(root_name_comps + ["backends", "rewrite"])
      )
  
      new_test_case_names = []
      for backend in backends:
          new_module_name_comps = (
              root_name_comps
              + ["dynamic", "backend_{}".format(backend)]
              + relative_module_name_comps
          )
          # Rewrite the module.
          new_module = importlib.import_module(
              ".".join(new_module_name_comps)
          )
  
          # Subclass the test case so that we can rename it (absl uses the class name
          # in its UI).
          base_new_test = getattr(new_module, test_case.__name__)
          new_test = type(
              "{}_{}".format(test_case.__name__, backend),
              (base_new_test,),
              {},
          )
          new_test_case_names.append(new_test.__name__)
          globals_dict[new_test.__name__] = new_test
  
imports: import collections
  import importlib
  from absl import flags
  from absl import logging
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling a global function with example usage: multi_backend_test(globals_dict, relative_module_name, backends, test_case) and returns: function, test_case