code: def test_gluon_dirichlet():
      class TestDirichlet(HybridBlock):
          def __init__(self, func):
              super(TestDirichlet, self).__init__()
              self._func = func
  
          def forward(self, alpha, *args):
              dirichlet = mgp.Dirichlet(alpha, validate_args=True)
              return _distribution_method_invoker(
                  dirichlet, self._func, *args
              )
  
      event_shapes = [2, 4, 6]
      batch_shapes = [None, (2, 3)]
  
      # Test sampling
      for event_shape, batch_shape in itertools.product(
          event_shapes, batch_shapes
      ):
          for hybridize in [True, False]:
              desired_shape = (
                  batch_shape if batch_shape is not None else ()
              ) + (event_shape,)
              alpha = np.random.uniform(1.0, 5.0, size=desired_shape)
              net = TestDirichlet("sample")
              if hybridize:
                  net.hybridize()
              mx_out = net(alpha).asnumpy()
              # Check shape
              assert mx_out.shape == desired_shape
              # Check simplex
              assert_almost_equal(
                  mx_out.sum(-1),
                  _np.ones_like(mx_out.sum(-1)),
                  atol=1e-4,
                  rtol=1e-3,
                  use_broadcast=False,
              )
  
      # Test log_prob
      # Scipy does not support batch `alpha`, thus we skip multi-dimensional batch_shape case.
      for event_shape, batch_shape in itertools.product(
          event_shapes, batch_shapes[:1]
      ):
          for hybridize in [True, False]:
              desired_shape = (
                  batch_shape if batch_shape is not None else ()
              ) + (event_shape,)
              alpha = np.random.uniform(1.0, 5.0, size=desired_shape)
              np_samples = _np.random.dirichlet(
                  [10.0 / event_shape] * event_shape, size=batch_shape
              )
              net = TestDirichlet("log_prob")
              if hybridize:
                  net.hybridize()
              mx_out = net(alpha, np.array(np_samples)).asnumpy()
              np_out = ss.dirichlet(alpha=alpha.asnumpy()).logpdf(
                  np_samples
              )
              assert_almost_equal(
                  mx_out,
                  np_out,
                  atol=1e-4,
                  rtol=1e-3,
                  use_broadcast=False,
              )
  
      # Test `mean`, `var` and `entropy`
      for event_shape, batch_shape in itertools.product(
          event_shapes, batch_shapes[:1]
      ):
          for hybridize in [True, False]:
              for func in ["mean", "variance", "entropy"]:
                  desired_shape = (
                      batch_shape if batch_shape is not None else ()
                  ) + (event_shape,)
                  alpha = np.random.uniform(1.0, 5.0, desired_shape)
                  net = TestDirichlet(func)
                  if hybridize:
                      net.hybridize()
                  mx_out = net(alpha).asnumpy()
                  ss_dir = ss.dirichlet(alpha=alpha.asnumpy())
                  if func == "mean":
                      np_out = ss_dir.mean()
                  elif func == "variance":
                      np_out = ss_dir.var()
                  else:
                      np_out = ss_dir.entropy()
                  assert_almost_equal(
                      mx_out,
                      np_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 15 for calling a global function with example usage: test_gluon_dirichlet() and returns: _distribution_method_invoker