code: def test_gluon_mvn():
      class TestMVN(HybridBlock):
          def __init__(self, func, param_type):
              super(TestMVN, self).__init__()
              self._func = func
              # cov, precision or scale_tril
              self._param_type = param_type
  
          def forward(self, loc, cov, *args):
              mvn = mgp.MultivariateNormal(
                  loc=loc, **{self._param_type: cov}, validate_args=True
              )
              return _distribution_method_invoker(
                  mvn, self._func, *args
              )
  
      def _stable_inv(cov):
          """
          Force the precision matrix to be symmetric.
          """
          precision = np.linalg.inv(cov)
          return (precision + precision.mT) / 2
  
      event_shapes = [3, 5]
      loc_shapes = [(), (2,), (4, 2)]
      cov_shapes = [(), (2,), (4, 2)]
      cov_func = {
          "cov": lambda s: s,
          "precision": lambda s: _stable_inv(s),
          "scale_tril": lambda s: np.linalg.cholesky(s),
      }
  
      # Test sampling
      for loc_shape, cov_shape, event_shape in itertools.product(
          loc_shapes, cov_shapes, event_shapes
      ):
          for cov_type in cov_func.keys():
              for hybridize in [True, False]:
                  loc = np.random.randn(*(loc_shape + (event_shape,)))
                  _s = np.random.randn(
                      *(cov_shape + (event_shape, event_shape))
                  )
                  loc.attach_grad()
                  _s.attach_grad()
                  # Full covariance matrix
                  sigma = np.matmul(_s, _s.mT) + np.eye(event_shape)
                  cov_param = cov_func[cov_type](sigma)
                  net = TestMVN("sample", cov_type)
                  if hybridize:
                      net.hybridize()
                  with autograd.record():
                      mx_out = net(loc, cov_param)
                  desired_shape = (loc + sigma[..., 0]).shape
                  assert mx_out.shape == desired_shape
                  mx_out.backward()
                  assert loc.grad.shape == loc.shape
                  assert _s.grad.shape == _s.shape
  
      # Test log_prob
      for loc_shape, cov_shape, event_shape in itertools.product(
          loc_shapes, cov_shapes, event_shapes
      ):
          for cov_type in cov_func.keys():
              for hybridize in [True, False]:
                  loc = np.random.randn(*(loc_shape + (event_shape,)))
                  _s = np.random.randn(
                      *(cov_shape + (event_shape, event_shape))
                  )
                  samples = np.random.normal(
                      np.zeros_like(loc), np.ones_like(_s[..., 0])
                  )
                  loc.attach_grad()
                  _s.attach_grad()
                  # Full covariance matrix
                  sigma = np.matmul(_s, _s.mT) + np.eye(event_shape)
                  cov_param = cov_func[cov_type](sigma)
                  net = TestMVN("log_prob", cov_type)
                  if hybridize:
                      net.hybridize()
                  mx_out = net(loc, cov_param, samples)
                  assert mx_out.shape == samples.shape[:-1]
                  if mx_out.shape == ():
                      mx_out_t = mx_out.asnumpy()
                  else:
                      mx_out_t = mx_out.asnumpy().flatten()[0]
                  # Select the first element in the batch, because scipy does not support batching.
                  loc_t = loc.reshape(-1, event_shape)[0].asnumpy()
                  sigma_t = sigma.reshape(-1, event_shape, event_shape)[
                      0
                  ].asnumpy()
                  samples_t = samples.reshape(
                      -1, event_shape
                  ).asnumpy()[0]
                  scipy_mvn = ss.multivariate_normal(loc_t, sigma_t)
                  ss_out = scipy_mvn.logpdf(samples_t)
                  assert_almost_equal(
                      mx_out_t,
                      ss_out,
                      atol=1e-4,
                      rtol=1e-3,
                      use_broadcast=False,
                  )
  
      # Test entropy
      for loc_shape, cov_shape, event_shape in itertools.product(
          loc_shapes, cov_shapes, event_shapes
      ):
          for cov_type in cov_func.keys():
              for hybridize in [True, False]:
                  loc = np.random.randn(*(loc_shape + (event_shape,)))
                  _s = np.random.randn(
                      *(cov_shape + (event_shape, event_shape))
                  )
                  loc.attach_grad()
                  _s.attach_grad()
                  # Full covariance matrix
                  sigma = np.matmul(_s, _s.mT) + np.eye(event_shape)
                  cov_param = cov_func[cov_type](sigma)
                  net = TestMVN("entropy", cov_type)
                  if hybridize:
                      net.hybridize()
                  mx_out = net(loc, cov_param)
                  assert mx_out.shape == sigma.shape[:-2]
                  if mx_out.shape == ():
                      mx_out_t = mx_out.asnumpy()
                  else:
                      mx_out_t = mx_out.asnumpy().flatten()[0]
                  # Select the first element in the batch, because scipy does not support batching.
                  loc_t = loc.reshape(-1, event_shape)[0].asnumpy()
                  sigma_t = sigma.reshape(-1, event_shape, event_shape)[
                      0
                  ].asnumpy()
                  scipy_mvn = ss.multivariate_normal(loc_t, sigma_t)
                  ss_out = scipy_mvn.entropy()
                  assert_almost_equal(
                      mx_out_t,
                      ss_out,
                      atol=1e-4,
                      rtol=1e-3,
                      use_broadcast=False,
                  )
  
imports: import mxnet
  import numpy
  from mxnet import np, npx, autograd
  from mxnet import gluon
  import mxnet.gluon.probability
  from mxnet.gluon.probability import StochasticBlock, StochasticSequential
  from mxnet.gluon import HybridBlock
  from mxnet.test_utils import use_np, assert_almost_equal
  from numpy.testing import assert_array_equal
  import pytest
  import scipy.stats
  import scipy.special
  import itertools
  from numbers import Number
label_desc: Use image node_id 27 for calling a global function with example usage: test_gluon_mvn() and returns: unknown, _distribution_method_invoker