class_docstr: Generates a tfp.distribution.Categorical by predicting logits.
code: def call(self, inputs, outer_rank, training=False, mask=None):
      # outer_rank is needed because the projection is not done on the raw
      # observations so getting the outer rank is hard as there is no spec to
      # compare to.
      batch_squash = utils.BatchSquash(outer_rank)
      inputs = batch_squash.flatten(inputs)
      inputs = tf.cast(inputs, tf.float32)
  
      logits = self._projection_layer(inputs, training=training)
      logits = tf.reshape(logits, [-1] + self._output_shape.as_list())
      logits = batch_squash.unflatten(logits)
  
      if mask is not None:
          # If the action spec says each action should be shaped (1,), add another
          # dimension so the final shape is (B, 1, A), where A is the number of
          # actions. This will make Categorical emit events shaped (B, 1) rather
          # than (B,). Using axis -2 to allow for (B, T, 1, A) shaped q_values.
          if mask.shape.rank < logits.shape.rank:
              mask = tf.expand_dims(mask, -2)
  
          # Overwrite the logits for invalid actions to a very large negative
          # number. We do not use -inf because it produces NaNs in many tfp
          # functions.
          almost_neg_inf = tf.constant(
              logits.dtype.min, dtype=logits.dtype
          )
          logits = tf.compat.v2.where(
              tf.cast(mask, tf.bool), logits, almost_neg_inf
          )
  
      return self.output_spec.build_distribution(logits=logits), ()
  
file_path: agents/tf_agents/networks/categorical_projection_network.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.specs import distribution_spec
  from tf_agents.specs import tensor_spec
label_desc: Use image node_id 3 for calling the CategoricalProjectionNetwork obj's underlying member method code with example usage: obj.call(inputs, outer_rank, training, mask) and returns: self,  