code: def test_categorical_dist_batch():
      n_classes = 2
      probs = mx.nd.array(((0.8, 0.2), (0.7, 0.3), (0.5, 0.5)))
  
      dist = CategoricalDist(n_classes, probs)
      # check log_prob
      actions = mx.nd.array((0, 1, 0))
      manual_log_prob = np.array(
          (-0.22314353, -1.20397282, -0.69314718)
      )
      np.testing.assert_almost_equal(
          actual=dist.log_prob(actions).asnumpy(),
          desired=manual_log_prob,
      )
      # check entropy
      sp_entropy = np.array(
          [
              sp_stats.entropy(pk=(0.8, 0.2)),
              sp_stats.entropy(pk=(0.7, 0.3)),
              sp_stats.entropy(pk=(0.5, 0.5)),
          ]
      )
      np.testing.assert_almost_equal(
          actual=dist.entropy().asnumpy(), desired=sp_entropy
      )
  
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 9 for calling a global function with example usage: test_categorical_dist_batch() without return types