code: def test_nanmixture_output(distribution_output, serialize_fn) -> None:
      nmdo = NanMixtureOutput(distribution_output)
  
      args_proj = nmdo.get_args_proj()
      args_proj.initialize()
  
      input = mx.nd.ones(shape=(3, 2))
  
      distr_args = args_proj(input)
  
      d = nmdo.distribution(distr_args)
      d = serialize_fn(d)
  
      samples = d.sample(num_samples=NUM_SAMPLES)
  
      sample = d.sample()
  
      assert samples.shape == (NUM_SAMPLES, *sample.shape)
  
      log_prob = d.log_prob(sample)
  
      assert log_prob.shape == d.batch_shape
  
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 5 for calling a global function with example usage: test_nanmixture_output(distribution_output, serialize_fn) without return types