class_docstr: Represents a parameterized discrete valued distribution.
  
  The returned distribution is essentially a `tfd.Categorical` that knows its
  support and thus can compute the mean value.
class_docstr_tok: ["Represents","a","parameterized","discrete","valued","distribution",".","The","returned","distribution","is","essentially","a","`","tfd.Categorical","`","that","knows","its","support","and","thus","can","compute","the","mean","value","."]
docstr: Initialization.
  
  If vmin and vmax have shape S, this will store the category values as a
  Tensor of shape (S*, num_atoms).
  
  Args:
    vmin: Minimum of the value range
    vmax: Maximum of the value range
    num_atoms: The atom values associated with each bin.
    logits_shape: The shape of the logits, excluding batch and num_atoms
      dimensions.
    w_init: Initialization for linear layer weights.
    b_init: Initialization for linear layer biases.
docstr_tok: ["Initialization",".","If","vmin","and","vmax","have","shape","S",",","this","will","store","the","category","values","as","a","Tensor","of","shape","(","S","*",",","num_atoms",")",".","Args",":","vmin",":","Minimum","of","the","value","range","vmax",":","Maximum","of","the","value","range","num_atoms",":","The","atom","values","associated","with","each","bin",".","logits_shape",":","The","shape","of","the","logits",",","excluding","batch","and","num_atoms","dimensions",".","w_init",":","Initialization","for","linear","layer","weights",".","b_init",":","Initialization","for","linear","layer","biases","."]
code: def __init__(
      self,
      vmin: float,
      vmax: float,
      num_atoms: int,
      logits_shape: Optional[Sequence[int]] = None,
      w_init: Optional[Initializer] = None,
      b_init: Optional[Initializer] = None,
  ):
      """Initialization.
  
      If vmin and vmax have shape S, this will store the category values as a
      Tensor of shape (S*, num_atoms).
  
      Args:
        vmin: Minimum of the value range
        vmax: Maximum of the value range
        num_atoms: The atom values associated with each bin.
        logits_shape: The shape of the logits, excluding batch and num_atoms
          dimensions.
        w_init: Initialization for linear layer weights.
        b_init: Initialization for linear layer biases.
      """
      super().__init__(name="DiscreteValuedHead")
      self._values = np.linspace(vmin, vmax, num=num_atoms, axis=-1)
      if not logits_shape:
          logits_shape = ()
      self._logits_shape = logits_shape + (num_atoms,)
      self._w_init = w_init
      self._b_init = b_init
  
imports: from typing import Any, List, Optional, Sequence, Union, Callable
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 1 to create a new DiscreteValuedTfpHead object from inherited base classes: hk with example: obj = DiscreteValuedTfpHead(vmin, vmax, num_atoms, logits_shape, w_init, b_init)