class_docstr: Module that produces a categorical distribution with the given number of values.
class_docstr_tok: ["Module","that","produces","a","categorical","distribution","with","the","given","number","of","values","."]
code: def __init__(
      self,
      num_values: Union[int, List[int]],
      dtype: Optional[Any] = jnp.int32,
      w_init: Optional[Initializer] = None,
      name: Optional[str] = None,
  ):
      super().__init__(name=name)
      self._dtype = dtype
      self._logit_shape = num_values
      self._linear = hk.Linear(np.prod(num_values), w_init=w_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 CategoricalHead object from inherited base classes: hk with example: obj = CategoricalHead(num_values, dtype, w_init, name)