class_docstr: Sets up distributed devices and sharding.
code: def setUp(self):
      super(DistributedTest, self).setUp()
      distributed_set_up()
      if JAX_MODE:
          self.key = jax.random.PRNGKey(0)
          assert jax.device_count() == NUM_DEVICES
      elif NUMPY_MODE:
          pass
      else:
          # Each test case needs a separate key base to not have the collective
          # clash.
          tf.distribute.MirroredStrategy._collective_key_base += 1
          # We need to defer creating the strategy, because some ops don't get
          # registered if we create it in setUp (which is ran eagerly) and then
          # used in TF1 graph mode.
          self._strategy = None
          self.key = [0, 0]
      self.axis_name = "i"
  
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 the DistributedTest obj's underlying member method code with example usage: obj.setUp() without return types