code: def test_categorical_dist_batch_time_shape():
      batch_size = 3
      time_steps = 4
      num_actions = 2
      # actions taken, of shape (batch_size, time_steps)
      actions = mx.nd.array(((0, 1, 0, 0), (1, 1, 0, 0), (0, 0, 0, 0)))
      # action probabilities, of shape (batch_size, time_steps, num_actions)
      policy_probs = mx.nd.array(
          (
              ((0.8, 0.2), (0.5, 0.5), (0.5, 0.5), (0.5, 0.5)),
              ((0.5, 0.5), (0.5, 0.5), (0.5, 0.5), (0.5, 0.5)),
              ((0.5, 0.5), (0.5, 0.5), (0.5, 0.5), (0.5, 0.5)),
          )
      )
      policy_dist = CategoricalDist(num_actions, policy_probs)
      action_probs = policy_dist.log_prob(actions)
      assert action_probs.ndim == 2
      assert action_probs.shape[0] == batch_size
      assert action_probs.shape[1] == time_steps
  
imports: import mxnet
  import numpy
  import os
  import pytest
  from scipy import stats
  import sys
  from rl_coach.architectures.head_parameters import PPOHeadParameters
  from rl_coach.architectures.mxnet_components.heads.ppo_head import CategoricalDist, MultivariateNormalDist, DiscretePPOHead, ClippedPPOLossDiscrete, ClippedPPOLossContinuous, PPOHead
  from rl_coach.agents.clipped_ppo_agent import ClippedPPOAlgorithmParameters, ClippedPPOAgentParameters
  from rl_coach.spaces import SpacesDefinition, DiscreteActionSpace
label_desc: Use image node_id 8 for calling a global function with example usage: test_categorical_dist_batch_time_shape() without return types