class_docstr: Sets up distributed devices and sharding.
code: def strategy_run(self, f, args=(), in_axes=0, axis_name=None):
      if JAX_MODE:
          axis_name = axis_name or self.axis_name
          if in_axes is None:
              return jax.pmap(
                  lambda _, args: f(*args),
                  in_axes=(0, None),
                  axis_name=axis_name,
              )(tf.ones(NUM_DEVICES), args)
          return jax.pmap(f, axis_name=axis_name, in_axes=in_axes)(
              *args
          )
      return self.strategy().run(tf.function(f, autograph=False), args)
  
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 5 for calling the DistributedTest obj's underlying member method code with example usage: obj.strategy_run(f, args, in_axes, axis_name) and returns: self, jax, jax