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","."]
code: def __call__(self, inputs: chex.Array) -> tfd.Distribution:
      net = hk.Linear(
          np.prod(self._logits_shape),
          w_init=self._w_init,
          b_init=self._b_init,
      )
      logits = net(inputs)
      logits = hk.Reshape(self._logits_shape, preserve_dims=1)(logits)
      return DiscreteValuedTfpDistribution(
          values=self._values, logits=logits
      )
  
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 2 for calling the DiscreteValuedTfpHead obj's underlying member method code with example usage: obj.__call__(inputs) and returns: DiscreteValuedTfpDistribution