code: def _external_rewrites(backend, data):
      """Rewrite external imports."""
      # TFP
      if backend == "backend_tensorflow":
          sub = "import tensorflow_probability as tfp"
      elif backend == "backend_jax":
          sub = "import tensorflow_probability.substrates.jax as tfp"
      elif backend == "backend_numpy":
          sub = "import tensorflow_probability.substrates.numpy as tfp"
      else:
          raise ValueError(f"Unknown backend {backend}")
      data = re.sub("import tensorflow_probability as tfp", sub, data)
  
      # TensorFlow
      if backend == "backend_tensorflow":
          sub = "import tensorflow.compat.v2 as tf"
      elif backend == "backend_jax":
          sub = (
              "from tensorflow_probability.substrates import jax as _tfp_jax; tf ="
              " _tfp_jax.tf2jax; del _tfp_jax"
          )
      elif backend == "backend_numpy":
          sub = (
              "from tensorflow_probability.substrates import numpy as _tfp_numpy; "
              "tf = _tfp_numpy.tf2numpy; del _tfp_numpy"
          )
      else:
          raise ValueError(f"Unknown backend {backend}")
      data = re.sub("import tensorflow.compat.v2 as tf", sub, data)
  
      # TFP's direct imports
      if backend == "backend_tensorflow":
          sub = "from tensorflow_probability.python"
      elif backend == "backend_jax":
          sub = "from tensorflow_probability.substrates.jax"
      elif backend == "backend_numpy":
          sub = "from tensorflow_probability.substrates.numpy"
      else:
          raise ValueError(f"Unknown backend {backend}")
      data = re.sub("from tensorflow_probability.python", sub, data)
  
      # TF's internal nest module
      if backend == "backend_tensorflow":
          sub = "from tensorflow.python.util import nest"
      elif backend == "backend_jax":
          sub = "from tensorflow_probability.python.internal.backend.jax import nest"
      elif backend == "backend_numpy":
          sub = "from tensorflow_probability.python.internal.backend.numpy import nest"
      else:
          raise ValueError(f"Unknown backend {backend}")
      data = re.sub(
          "from tensorflow.python.util import nest", sub, data
      )
  
      return data
  
file_path: probability/spinoffs/inference_gym/inference_gym/backends/rewrite.py
imports: import importlib
  import importlib.abc
  import importlib.machinery
  import importlib.util
  import re
  import sys
label_desc: Use image node_id 1 for calling a global function with example usage: _external_rewrites(backend, data) and returns: data