class_docstr: Sets up distributed devices and sharding.
code: def shard_values(self, values, axis=0):
      self.assertEqual(values.shape[axis], NUM_DEVICES)
  
      if JAX_MODE:
          return values
  
      values = distribution_util.move_dimension(values, axis, 0)
  
      def value_fn(ctx):
          return values[ctx.replica_id_in_sync_group]
  
      return (
          self.strategy().experimental_distribute_values_from_function(
              value_fn
          )
      )
  
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 6 for calling the DistributedTest obj's underlying member method code with example usage: obj.shard_values(values, axis) and returns: self, values, values