code: def distributed_set_up():
      """Sets up virtual devices to be used for testing distributed code."""
      if JAX_MODE:
          os.environ["XLA_FLAGS"] = (
              f'{os.environ.get("XLA_FLAGS", "")} '
              f"--xla_force_host_platform_device_count={NUM_DEVICES}"
          )
      elif NUMPY_MODE:
          return
      else:
          physical_devices = (
              tf.config.experimental.list_physical_devices()
          )
  
          tf.config.experimental.set_virtual_device_configuration(
              physical_devices[0],
              [tf.config.experimental.VirtualDeviceConfiguration()]
              * NUM_DEVICES,
          )
  
file_path: probability/tensorflow_probability/python/internal/distribute_test_lib.py
imports: import os
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling a global function with example usage: distributed_set_up() without return types