class_docstr: A meta-bijector meant for use in an SPMD distributed context.
  
  `Sharded` is a meta-bijector that enables distributions to be used in SPMD
  programs. A `Sharded` bijector represents a bijector which acts on elements
  are split across a set of devices. The number of shards is the number of
  devices in the current TensorFlow DistributionStrategy or the provided JAX
  pmap axis.
docstr: Constructs a `Sharded` bijector.
  
  Args:
    bijector: The base bijector instance to transform. Typically an instance
      of `Bijector`.
    shard_axis_name: `str` for axis name for use in JAX backend.
    name: The name for ops managed by the distribution.
      Default value: `None` (i.e., `'Sharded' + bijector.name`).
code: def __init__(self, bijector, *, shard_axis_name, name=None):
      """Constructs a `Sharded` bijector.
  
      Args:
        bijector: The base bijector instance to transform. Typically an instance
          of `Bijector`.
        shard_axis_name: `str` for axis name for use in JAX backend.
        name: The name for ops managed by the distribution.
          Default value: `None` (i.e., `'Sharded' + bijector.name`).
      """
      if JAX_MODE and shard_axis_name is None:
          raise ValueError(
              "Cannot provide a `None` axis name in JAX backend."
          )
  
      with tf.name_scope(name or "Sharded" + bijector.name) as name:
          super().__init__(
              name=name,
              forward_min_event_ndims=bijector.forward_min_event_ndims,
              inverse_min_event_ndims=bijector.inverse_min_event_ndims,
          )
          self._bijector = bijector
          self._shard_axis_name = shard_axis_name
  
file_path: probability/tensorflow_probability/python/experimental/bijectors/sharded.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.internal import distribute_lib
label_desc: Use image node_id 1 to create a new Sharded object from inherited base classes: bijector_lib with example: obj = Sharded(bijector)