code: def test_nan_mixture(
      distr_class,
      p: Tensor,
      x: Tensor,
      distr_params: Dict[str, Tensor],
      distr_params_grad: Dict[str, Tensor],
      serialize_fn,
  ) -> None:
      # sample from component distributions, and select samples
      distr = distr_class(**distr_params)
  
      samples = distr.sample(num_samples=NUM_SAMPLES_LARGE)
  
      rand = mx.nd.random.uniform(shape=(NUM_SAMPLES_LARGE, *p.shape))
      choice = (rand > p.expand_dims(axis=0)).broadcast_like(samples)
      samples_ref = mx.nd.where(choice, samples, samples.zeros_like())
  
      # construct NanMixture distribution and sample from it
      nan_mixture = NanMixture(nan_prob=p, distribution=distr)
  
      nan_mixture = serialize_fn(nan_mixture)
  
      samples_mix = nan_mixture.sample(num_samples=NUM_SAMPLES_LARGE)
      # check that shapes are right
  
      assert samples.shape == samples_mix.shape == samples_ref.shape
  
      # TODO check mean and stddev
  
      # check log_prob
      log_prob = nan_mixture.log_prob(x)
  
      log_prob_true = mx.nd.log(
          mx.nd.where(x != x, p, (1 - p) * distr.prob(x))
      )
  
      assert np.allclose(log_prob.asnumpy(), log_prob_true.asnumpy())
  
      for param in distr_params:
          distr_params[param].attach_grad()
      p.attach_grad()
  
      with mx.autograd.record():
          distr = distr_class(**distr_params)
          nan_mixture = NanMixture(nan_prob=p, distribution=distr)
          nll = -nan_mixture.log_prob(x)
      nll.backward()
  
      p_grad_true = mx.nd.where(x != x, -1 / p, 1 / (1 - p))
      # gradient is undefined for these cases:
      p_grad_true = mx.nd.where(
          mx.nd.logical_or(
              mx.nd.logical_and(x != x, p == 0),
              mx.nd.logical_and(x == x, p == 1),
          ),
          0.0 / p_grad_true.zeros_like(),
          p_grad_true,
      )
  
      assert np.allclose(p.grad.asnumpy(), p_grad_true.asnumpy())
  
      for param in distr_params:
          assert np.allclose(
              distr_params[param].grad.asnumpy(),
              distr_params_grad[param],
          )
  
file_path: gluonts/test/mx/distribution/test_nan_mixture.py
imports: from typing import Dict
  import mxnet
  import numpy
  import pytest
  from gluonts.core.serde import dump_json, load_json
  from gluonts.gluonts_tqdm import tqdm
  from gluonts.mx import Tensor
  from gluonts.mx.distribution import Categorical, CategoricalOutput, Gaussian, GaussianOutput, NanMixture, NanMixtureOutput, StudentTOutput
label_desc: Use image node_id 2 for calling a global function with example usage: test_nan_mixture(distr_class, p, x, distr_params, distr_params_grad, serialize_fn) without return types