code: def test_distribution_output_shapes(
      distr_out: DistributionOutput,
      data: Tensor,
      loc: List[Union[None, Tensor]],
      scale: List[Union[None, Tensor]],
      expected_batch_shape: Tuple,
      expected_event_shape: Tuple,
  ):
      args_proj = distr_out.get_args_proj()
      args_proj.initialize()
  
      args = args_proj(data)
  
      assert distr_out.event_shape == expected_event_shape
  
      for l, s in product(loc, scale):
          distr = distr_out.distribution(args, loc=l, scale=s)
  
          assert distr.batch_shape == expected_batch_shape
          assert distr.event_shape == expected_event_shape
  
          x = distr.sample()
  
          assert x.shape == distr.batch_shape + distr.event_shape
  
          loss = distr.loss(x)
  
          assert loss.shape == distr.batch_shape
  
          x1 = distr.sample(num_samples=1)
  
          assert (
              x1.shape == (1,) + distr.batch_shape + distr.event_shape
          )
  
          x3 = distr.sample(num_samples=3)
  
          assert (
              x3.shape == (3,) + distr.batch_shape + distr.event_shape
          )
  
file_path: gluonts/test/mx/distribution/test_distribution_output_shapes.py
imports: import pytest
  from itertools import product
  from typing import List, Tuple, Union
  import mxnet
  from gluonts.mx import Tensor
  from gluonts.mx.distribution import BetaOutput, DeterministicOutput, DirichletMultinomialOutput, DirichletOutput, DistributionOutput, GammaOutput, GaussianOutput, LaplaceOutput, LowrankMultivariateGaussianOutput, MixtureDistributionOutput, MultivariateGaussianOutput, NegativeBinomialOutput, PiecewiseLinearOutput, PoissonOutput, StudentTOutput, UniformOutput
label_desc: Use image node_id 1 for calling a global function with example usage: test_distribution_output_shapes(distr_out, data, loc, scale, expected_batch_shape, expected_event_shape) without return types