class_docstr: Represents a parameterized discrete valued distribution.
  
  The returned distribution is essentially a `tfd.Categorical`, but one which
  knows its support and so 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.
    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",".","w_init",":","Initialization","for","linear","layer","weights",".","b_init",":","Initialization","for","linear","layer","biases","."]
code: def __init__(
      self,
      vmin: Union[float, np.ndarray, tf.Tensor],
      vmax: Union[float, np.ndarray, tf.Tensor],
      num_atoms: int,
      w_init: Optional[snt.initializers.Initializer] = None,
      b_init: Optional[snt.initializers.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.
        w_init: Initialization for linear layer weights.
        b_init: Initialization for linear layer biases.
      """
      super().__init__(name="DiscreteValuedHead")
      vmin = tf.convert_to_tensor(vmin)
      vmax = tf.convert_to_tensor(vmax)
      self._values = tf.linspace(vmin, vmax, num_atoms, axis=-1)
      self._distributional_layer = snt.Linear(
          tf.size(self._values), w_init=w_init, b_init=b_init
      )
  
imports: import types
  from typing import Optional, Union
  from absl import logging
  from acme.tf.networks import distributions
  import numpy
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 1 to create a new DiscreteValuedHead object from inherited base classes: snt with example: obj = DiscreteValuedHead(vmin, vmax, num_atoms, w_init, b_init)