code: def test_gluon_stochastic_block():
      class dummyBlock(StochasticBlock):
          """In this test case, we generate samples from a Gaussian parameterized
          by `loc` and `scale` and accumulate the KL-divergence between it and
          its prior and the l2 norm of `loc` into the block's loss storage.
          """
  
          @StochasticBlock.collectLoss
          def forward(self, loc, scale):
              qz = mgp.Normal(loc, scale)
              # prior
              pz = mgp.Normal(np.zeros_like(loc), np.ones_like(scale))
              self.add_loss(mgp.kl_divergence(qz, pz))
              self.add_loss((loc**2).sum(1))
              return qz.sample()
  
      shape = (4, 4)
      for hybridize in [True, False]:
          net = dummyBlock()
          if hybridize:
              net.hybridize()
          loc = np.random.randn(*shape)
          scale = np.random.rand(*shape)
          mx_out = net(loc, scale).asnumpy()
          kl = net.losses[0].asnumpy()
          l2_norm = net.losses[1].asnumpy()
          assert mx_out.shape == loc.shape
          assert kl.shape == loc.shape
          assert l2_norm.shape == shape[:-1]
          if hybridize:
              net.export("dummyBlock", epoch=0)
  
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 34 for calling a global function with example usage: test_gluon_stochastic_block() and returns: qz