code: def test_gluon_kl():
      def _test_zero_kl(p, shape):
          """Check if KL(p || p) = 0
  
          Parameters
          ----------
          p : Distribution
          """
          mx_out = mgp.kl_divergence(p, p).asnumpy()
          np_out = _np.zeros(shape)
          assert_almost_equal(
              mx_out, np_out, atol=1e-3, rtol=1e-2, use_broadcast=False
          )
  
      def _test_monte_carlo(p, q, M=50000):
          r"""Check if KL(p || q) is approximately equal to
          1/M * \Sum_{i=1}^{M} log(p(x_i) / q(x_i)), x_i ~ p(x)
          """
          kl = mgp.kl_divergence(p, q)
          mc_approx = mgp.empirical_kl(p, q, M)
          assert_almost_equal(
              mc_approx.asnumpy(),
              kl.asnumpy(),
              atol=1e-1,
              rtol=1e-1,
              use_broadcast=False,
          )
  
      def _dist_factory(dist, *param_funcs):
          """Generate a distribution object with parameters of random value.
  
          Parameters
          ----------
          dist : Type
              A type of distribution.
          param_funcs : List
              A list of functions that generate valid parameters for `dist`
          """
          params = [f() if callable(f) else f for f in param_funcs]
          return dist(*params)
  
      # could cause longer runtime and potential flaky tests
      monte_carlo_test = False
      repeated_times = 50000
      shapes = [(), (1,), (2, 3), 6]
  
      # Test kl between same distributions
      # uniform
      for shape in shapes:
          dist = mgp.Uniform
  
          def low():
              return np.random.uniform(0, 1, shape)
  
          def high():
              return np.random.uniform(1, 2, shape)
  
          _test_zero_kl(_dist_factory(dist, low, high), shape)
  
      # normal, laplace, cauchy, gumbel
      for dist in [mgp.Normal, mgp.Laplace, mgp.Cauchy, mgp.Gumbel]:
          for shape in shapes:
  
              def loc():
                  return np.random.uniform(-1, 1, shape)
  
              def scale():
                  return np.random.uniform(0.5, 1.5, shape)
  
              _test_zero_kl(_dist_factory(dist, loc, scale), shape)
              if monte_carlo_test:
                  _test_monte_carlo(
                      _dist_factory(dist, loc, scale),
                      _dist_factory(dist, loc, scale),
                      repeated_times,
                  )
  
      # poisson
      for shape in shapes[1:]:
          dist = mgp.Poisson
  
          def rate():
              return np.random.uniform(0.5, 1.5, shape)
  
          _test_zero_kl(_dist_factory(dist, rate), shape)
          if monte_carlo_test:
              _test_monte_carlo(
                  _dist_factory(dist, rate),
                  _dist_factory(dist, rate),
                  repeated_times,
              )
  
      # exponential, geometric
      for dist in [mgp.Exponential, mgp.Geometric]:
          for shape in shapes:
  
              def s():
                  return np.random.uniform(size=shape, low=1e-3)
  
              _test_zero_kl(_dist_factory(dist, s), shape)
              if monte_carlo_test:
                  _test_monte_carlo(
                      _dist_factory(dist, s),
                      _dist_factory(dist, s),
                      repeated_times,
                  )
  
      # pareto
      for shape in shapes:
          dist = mgp.Pareto
  
          def alpha():
              return np.random.uniform(size=shape)
  
          def scale():
              return np.random.uniform(size=shape)
  
          _test_zero_kl(_dist_factory(dist, alpha, scale), shape)
  
      for shape in shapes:
          dist = mgp.HalfNormal
  
          def scale():
              return np.random.uniform(0.5, 1.5, shape)
  
          _test_zero_kl(_dist_factory(dist, scale), shape)
          if monte_carlo_test:
              _test_monte_carlo(
                  _dist_factory(dist, scale),
                  _dist_factory(dist, scale),
                  repeated_times,
              )
  
      # gamma, beta
      for dist in [mgp.Gamma, mgp.Beta]:
          for shape in shapes:
  
              def param1():
                  return np.random.uniform(0.5, 1.5, shape)
  
              def param2():
                  return np.random.uniform(0.5, 1.5, shape)
  
              _test_zero_kl(_dist_factory(dist, param1, param2), shape)
              if monte_carlo_test:
                  _test_monte_carlo(
                      _dist_factory(dist, param1, param2),
                      _dist_factory(dist, param1, param2),
                      50000,
                  )
  
      # binomial
      for shape in shapes:
          n = _np.random.randint(5, 10)
          prob = np.random.uniform(low=0.1, size=shape)
          dist = mgp.Binomial(n=n, prob=prob)
          _test_zero_kl(dist, shape)
  
      # bernoulli
      for shape in shapes:
          prob = np.random.uniform(size=shape)
          dist = mgp.Bernoulli(prob=prob)
          _test_zero_kl(dist, shape)
  
      event_shapes = [3, 5, 10]
      loc_shapes = [(), (2,), (4, 2)]
      cov_shapes = [(), (2,), (4, 2)]
      for loc_shape, cov_shape, event_shape in itertools.product(
          loc_shapes, cov_shapes, event_shapes
      ):
          loc = np.random.randn(*(loc_shape + (event_shape,)))
          _s = np.random.randn(
              *(cov_shape + (event_shape, event_shape))
          )
          sigma = np.matmul(_s, _s.mT) + np.eye(event_shape)
          dist = mgp.MultivariateNormal(loc, cov=sigma)
          desired_shape = (loc + sigma[..., 0]).shape[:-1]
          _test_zero_kl(dist, desired_shape)
  
      batch_shapes = loc_shapes
      # dirichlet
      for event_shape, batch_shape in itertools.product(
          event_shapes, batch_shapes
      ):
          desired_shape = batch_shape if batch_shape is not None else ()
          dist = mgp.Dirichlet
  
          def alpha():
              return np.random.uniform(
                  0.5, 1.5, size=(desired_shape + (event_shape,))
              )
  
          _test_zero_kl(_dist_factory(dist, alpha), desired_shape)
          if monte_carlo_test:
              _test_monte_carlo(
                  _dist_factory(dist, alpha),
                  _dist_factory(dist, alpha),
                  50000,
              )
  
      # categorical, One-hot categorical
      for dist in [mgp.Categorical, mgp.OneHotCategorical]:
          for event_shape, batch_shape in itertools.product(
              event_shapes, batch_shapes
          ):
              prob = lambda: np.array(
                  _np.random.dirichlet(
                      [1 / event_shape] * event_shape, size=batch_shape
                  )
              )
              _test_zero_kl(
                  _dist_factory(dist, event_shape, prob), batch_shape
              )
              if monte_carlo_test:
                  _test_monte_carlo(
                      _dist_factory(dist, event_shape, prob),
                      _dist_factory(dist, event_shape, prob),
                      repeated_times,
                  )
  
      # Test kl between different distributions
      # KL(Uniform || ...)
      for shape in shapes:
          rhs_dists = [
              mgp.Normal(
                  np.random.uniform(-1, 1, shape),
                  np.random.uniform(0.5, 1.5, shape),
              ),
              mgp.Gumbel(
                  np.random.uniform(-1, 1, shape),
                  np.random.uniform(0.5, 1.5, shape),
              ),
          ]
          for rhs_dist in rhs_dists:
              low = np.random.uniform(-1, 1, shape)
              high = low + np.random.uniform(0.5, 1.5, shape)
              lhs_dist = mgp.Uniform(low, high)
              kl = mgp.kl_divergence(lhs_dist, rhs_dist)
              assert kl.shape == low.shape
              if monte_carlo_test:
                  _test_monte_carlo(lhs_dist, rhs_dist, repeated_times)
  
      # KL(Exponential || ...)
      for shape in shapes:
          rhs_dists = [
              mgp.Normal(
                  np.random.uniform(-1, 1, shape),
                  np.random.uniform(0.5, 1.5, shape),
              ),
              mgp.Gumbel(
                  np.random.uniform(-1, 1, shape),
                  np.random.uniform(0.5, 1.5, shape),
              ),
              mgp.Gamma(
                  np.random.uniform(0.5, 1.5, shape),
                  np.random.uniform(0.5, 1.5, shape),
              ),
          ]
          for rhs_dist in rhs_dists:
              s = np.random.uniform(size=shape)
              lhs_dist = mgp.Exponential(s)
              kl = mgp.kl_divergence(lhs_dist, rhs_dist)
              assert kl.shape == s.shape
              if monte_carlo_test:
                  _test_monte_carlo(lhs_dist, rhs_dist, repeated_times)
  
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 33 for calling a global function with example usage: test_gluon_kl() and returns: dist, np, np, np, np, np, np, np, np, np, np, np, np