class_docstr: A meta-distribution meant for use in an SPMD distributed context.
  
  `Sharded` is a meta-distribution that enables distributions to be used in SPMD
  programs. A `Sharded` distribution represents a random variable that has
  been 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.
  
  In practice, `Sharded` modifies its input distribution in two ways.
  First, when a `Sharded` distribution is sampled, it first folds the current
  device ID into the input random seed, resulting in different samples on each
  device. Second, when computing the `log_prob` of a value, a `Sharded`
  distribution aggregates the log-prob over all devices, resulting in the same
  synchronized value.
class_docstr_tok: ["A","meta-distribution","meant","for","use","in","an","SPMD","distributed","context",".","`","Sharded","`","is","a","meta-distribution","that","enables","distributions","to","be","used","in","SPMD","programs",".","A","`","Sharded","`","distribution","represents","a","random","variable","that","has","been","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",".","In","practice",",","`","Sharded","`","modifies","its","input","distribution","in","two","ways",".","First",",","when","a","`","Sharded","`","distribution","is","sampled",",","it","first","folds","the","current","device","ID","into","the","input","random","seed",",","resulting","in","different","samples","on","each","device",".","Second",",","when","computing","the","`","log_prob","`","of","a","value",",","a","`","Sharded","`","distribution","aggregates","the","log-prob","over","all","devices",",","resulting","in","the","same","synchronized","value","."]
docstr: Constructs a `Sharded` distribution.
  
  Args:
    distribution: The base distribution instance to transform. Typically an
      instance of `Distribution`.
    shard_axis_name: `str` or a list of strings for axis name(s). An empty
      list means that no sharding is actually done. This can be `None` under
      the TensorFlow backend (meaning a sharded axis is present, but
      anonymous). Only the JAX backend supports multiple axes names.
    validate_args: Python `bool`.  Whether to validate input with asserts. If
      `validate_args` is `False`, and the inputs are invalid, correct behavior
      is not guaranteed.
    name: The name for ops managed by the distribution.
      Default value: `None` (i.e., `'Sharded' + distribution.name`).
docstr_tok: ["Constructs","a","`","Sharded","`","distribution",".","Args",":","distribution",":","The","base","distribution","instance","to","transform",".","Typically","an","instance","of","`","Distribution","`",".","shard_axis_name",":","`","str","`","or","a","list","of","strings","for","axis","name","(","s",")",".","An","empty","list","means","that","no","sharding","is","actually","done",".","This","can","be","`","None","`","under","the","TensorFlow","backend","(","meaning","a","sharded","axis","is","present",",","but","anonymous",")",".","Only","the","JAX","backend","supports","multiple","axes","names",".","validate_args",":","Python","`","bool","`",".","Whether","to","validate","input","with","asserts",".","If","`","validate_args","`","is","`","False","`",",","and","the","inputs","are","invalid",",","correct","behavior","is","not","guaranteed",".","name",":","The","name","for","ops","managed","by","the","distribution",".","Default","value",":","`","None","`","(","i.e.",",","`","'Sharded","'","+","distribution.name","`",")","."]
code: def __init__(
      self,
      distribution,
      shard_axis_name=None,
      validate_args=False,
      name=None,
  ):
      """Constructs a `Sharded` distribution.
  
      Args:
        distribution: The base distribution instance to transform. Typically an
          instance of `Distribution`.
        shard_axis_name: `str` or a list of strings for axis name(s). An empty
          list means that no sharding is actually done. This can be `None` under
          the TensorFlow backend (meaning a sharded axis is present, but
          anonymous). Only the JAX backend supports multiple axes names.
        validate_args: Python `bool`.  Whether to validate input with asserts. If
          `validate_args` is `False`, and the inputs are invalid, correct behavior
          is not guaranteed.
        name: The name for ops managed by the distribution.
          Default value: `None` (i.e., `'Sharded' + distribution.name`).
      """
      parameters = dict(locals())
  
      if not auto_composite_tensor.is_composite_tensor(distribution):
          raise ValueError(
              "`distribution` must be a `CompositeTensor`."
          )
  
      if shard_axis_name is None:
          if JAX_MODE:
              # In JAX, axes names matter and we don't know which axis name the user
              # might intend, so we bail.
              raise ValueError(
                  "Cannot provide a `None` axis name in JAX backend."
              )
          else:
              # In TF, there are no axes names, so we can pick a reasonable default.
              shard_axis_name = [True]
  
      # Use inner axes before outer axes
      full_shard_axis_name = (
          distribution.experimental_shard_axis_names
          + distribute_lib.canonicalize_named_axis(shard_axis_name)
      )
  
      if not JAX_MODE:
          if len(full_shard_axis_name) > 1:
              raise ValueError(
                  "TensorFlow backend does not support multiple shard axes:\n"
                  "inner shard_axis_names: "
                  f"{list(distribution.experimental_shard_axis_names)}\n"
                  f"outer shard_axis_names: {list(shard_axis_name)}"
              )
  
      if len(set(full_shard_axis_name)) != len(full_shard_axis_name):
          duplicates = set()
          seen = set()
          for axis_name in full_shard_axis_name:
              if axis_name in seen:
                  duplicates.add(axis_name)
              seen.add(axis_name)
          raise ValueError(
              "Found duplicate axis name(s).\n"
              "inner shard_axis_names: "
              f"{list(distribution.experimental_shard_axis_names)}\n"
              f"outer shard_axis_names: {shard_axis_name}\n"
              f"duplicates: {list(duplicates)}"
          )
  
      with tf.name_scope(name or "Sharded" + distribution.name) as name:
          self._distribution = distribution
          self._shard_axis_name = full_shard_axis_name
          super(Sharded, self).__init__(
              dtype=self._distribution.dtype,
              validate_args=validate_args,
              allow_nan_stats=self._distribution.allow_nan_stats,
              reparameterization_type=self._distribution.reparameterization_type,
              parameters=parameters,
              name=name,
          )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.experimental.bijectors import sharded
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 1 to create a new Sharded object from inherited base classes: distribution_lib with example: obj = Sharded(distribution, shard_axis_name, validate_args, name)