class_docstr: Sets up distributed devices and sharding.
code: def per_replica_to_tensor(self, value, axis=0):
      if JAX_MODE:
          # JAX, by default, stacks outputs along the first axis.
          if axis == 0:
              return value
          else:
              return tf.nest.map_structure(
                  lambda v: distribution_util.move_dimension(
                      v, 0, axis
                  ),
                  value,
              )
      return tf.nest.map_structure(
          lambda per_replica: tf.stack(per_replica.values, axis=axis),
          value,
      )
  
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 2 for calling the DistributedTest obj's underlying member method code with example usage: obj.per_replica_to_tensor(value, axis) and returns: tf, value, tf