class_docstr: C51-style head.
  
  For each action, it produces the logits for a discrete distribution over
  atoms. Therefore, the returned logits represents several distributions, one
  for each action.
class_docstr_tok: ["C51-style","head",".","For","each","action",",","it","produces","the","logits","for","a","discrete","distribution","over","atoms",".","Therefore",",","the","returned","logits","represents","several","distributions",",","one","for","each","action","."]
code: def __call__(self, inputs: jnp.ndarray):
      q_logits = self._network(inputs)
      q_logits = jnp.reshape(
          q_logits, (-1, self._num_actions, self._num_atoms)
      )
      q_dist = jax.nn.softmax(q_logits)
      q_values = jnp.sum(q_dist * self._atoms, axis=2)
      q_values = jax.lax.stop_gradient(q_values)
      return q_values, q_logits, self._atoms
  
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 DiscreteValued obj's underlying member method code with example usage: obj.__call__(inputs) and returns: q_values,  q_logits,  self