code: def _build_module(self, input_layer):
      values_distribution = self.dense_layer(
          self.num_actions * self.num_atoms
      )(
          input_layer,
          name="output",
          bias_initializer=self.output_bias_initializer,
      )
      values_distribution = tf.reshape(
          values_distribution,
          (
              tf.shape(values_distribution)[0],
              self.num_actions,
              self.num_atoms,
          ),
      )
      # softmax on atoms dimension
      self.output = tf.nn.softmax(values_distribution)
  
      # calculate cross entropy loss
      self.distributions = tf.placeholder(
          tf.float32,
          shape=(None, self.num_actions, self.num_atoms),
          name="distributions",
      )
      self.target = self.distributions
      self.loss = tf.nn.softmax_cross_entropy_with_logits(
          labels=self.target, logits=values_distribution
      )
      tf.losses.add_loss(self.loss)
  
      self.q_values = tf.tensordot(
          tf.cast(self.output, tf.float64), self.z_values, 1
      )
  
      # used in batch-rl to estimate a probablity distribution over actions
      self.softmax = self.add_softmax_with_temperature()
  
file_path: coach/rl_coach/architectures/tensorflow_components/heads/categorical_q_head.py
imports: import tensorflow
  import numpy
  from rl_coach.architectures.tensorflow_components.heads import QHead
  from rl_coach.architectures.tensorflow_components.layers import Dense
  from rl_coach.base_parameters import AgentParameters
  from rl_coach.spaces import SpacesDefinition
label_desc: Use image node_id 2 for calling the CategoricalQHead obj's underlying member method code with example usage: obj._build_module(input_layer) without return types