code: def multiaxis_categorical(  # pylint: disable=invalid-name
      q_tm1: networks.DiscreteValuedDistribution,
      r_t: tf.Tensor,
      d_t: tf.Tensor,
      q_t: networks.DiscreteValuedDistribution,
  ) -> tf.Tensor:
      """Implements a multi-axis categorical distributional TD(0)-learning loss.
  
      All arguments may have a leading batch axis, but q_tm1.logits, and one of
      r_t or d_t *must* have a leading batch axis.
  
      Args:
        q_tm1: Previous timestep's value distribution.
        r_t: Reward.
        d_t: Discount.
        q_t: Current timestep's value distribution.
  
      Returns:
        Cross-entropy Bellman loss between q_tm1 and q_t + r_t * d_t.
        Shape: (B, *E), where
          B is the batch size.
          E is the broadcasted shape of r_t, d_t, and q_t.values[:-1].
      """
      tf.assert_equal(tf.rank(r_t), tf.rank(d_t))
  
      # Append a singleton axis corresponding to the axis that indexes the atoms in
      # q_t.values.
      r_t = r_t[..., None]  # shape: (B, *R, 1)
      d_t = d_t[..., None]  # shape: (B, *D, 1)
  
      z_t = r_t + d_t * q_t.values  # shape: (B, *E, N)
  
      p_t = tf.nn.softmax(q_t.logits)
  
      # Performs L2 projection.
      target = tf.stop_gradient(
          multiaxis_l2_project(z_t, p_t, q_t.values)
      )
  
      # Calculates loss.
      loss = tf.nn.softmax_cross_entropy_with_logits(
          logits=q_tm1.logits, labels=target
      )
  
      return loss
  
imports: from acme.tf import networks
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: multiaxis_categorical(q_tm1, r_t, d_t, q_t) and returns: loss