class_docstr: Shards a transition kernel across a named axis.
  
  Ordinarily, one can produce independent Markov chains from a single kernel by
  proving a batch of states but when using named axes inside of a map (say
  in the case of using JAX's `pmap`, `vmap`, or `xmap`), the kernel is provided
  with state without batch dimensions. In order to sample independently across
  the named axis, the PRNG seed across the named axis must be different. This
  can be accomplished by folding the named axis index into the random seed.
  A `Sharded` kernel does exactly this, creating independent chains across a
  named axis.
docstr: Constructs a `Sharded` transition kernel.
  
  Args:
    inner_kernel: A `TransitionKernel` to be sharded.
    chain_axis_names: A `str` or list of `str`s that determine the named axes
      that independent Markov chains will be sharded across.
    validate_args: Python `bool`. When `True` kernel parameters are checked
      for validity. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.
docstr_tok: ["Constructs","a","`","Sharded","`","transition","kernel",".","Args",":","inner_kernel",":","A","`","TransitionKernel","`","to","be","sharded",".","chain_axis_names",":","A","`","str","`","or","list","of","`","str","`","s","that","determine","the","named","axes","that","independent","Markov","chains","will","be","sharded","across",".","validate_args",":","Python","`","bool","`",".","When","`","True","`","kernel","parameters","are","checked","for","validity",".","When","`","False","`","invalid","inputs","may","silently","render","incorrect","outputs",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class","."]
code: def __init__(
      self,
      inner_kernel,
      chain_axis_names,
      validate_args=False,
      name=None,
  ):
      """Constructs a `Sharded` transition kernel.
  
      Args:
        inner_kernel: A `TransitionKernel` to be sharded.
        chain_axis_names: A `str` or list of `str`s that determine the named axes
          that independent Markov chains will be sharded across.
        validate_args: Python `bool`. When `True` kernel parameters are checked
          for validity. When `False` invalid inputs may silently render incorrect
          outputs.
        name: Python `str` name prefixed to Ops created by this class.
      """
      chain_axis_names = distribute_lib.canonicalize_named_axis(
          chain_axis_names
      )
      self._parameters = dict(
          inner_kernel=inner_kernel,
          chain_axis_names=chain_axis_names,
          validate_args=validate_args,
          name=name,
      )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/sharded.py
imports: from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.mcmc import kernel
label_desc: Use image node_id 1 to create a new Sharded object from inherited base classes: kernel with example: obj = Sharded(inner_kernel, chain_axis_names, validate_args, name)