code: def test_distribution_output_mean(
      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)
  
      for l, s in product(loc, scale):
          distr = distr_out.distribution(args, loc=l, scale=s)
          assert (
              distr.mean.shape
              == expected_batch_shape + expected_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 2 for calling a global function with example usage: test_distribution_output_mean(distr_out, data, loc, scale, expected_batch_shape, expected_event_shape) without return types