code: def test_multivariate_normal_dist_batch_shape():
      num_var = 2
      batch_size = 3
      means = mx.nd.random.uniform(shape=(batch_size, num_var))
      # create batch of covariance matrices only defined on diagonal
      std = mx.nd.array((1, 0.5)).broadcast_like(means).expand_dims(-2)
      covar = mx.nd.eye(N=num_var) * std
      data = mx.nd.random.uniform(shape=(batch_size, num_var))
      policy_dist = MultivariateNormalDist(num_var, means, covar)
      log_probs = policy_dist.log_prob(data)
      assert log_probs.ndim == 1
      assert log_probs.shape[0] == batch_size
  
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 2 for calling a global function with example usage: test_multivariate_normal_dist_batch_shape() without return types