code: def test_lowrank_multivariate_gaussian(
      hybridize: bool, rank: int
  ) -> None:
      num_samples = 2000
      dim = 2
  
      mu = np.arange(0, dim) / float(dim)
      D = np.eye(dim) * (np.arange(dim) / dim + 0.5)
      if rank > 0:
          W = np.sqrt(np.ones((dim, rank)) * 0.2)
          Sigma = D + W.dot(W.transpose())
          W = mx.nd.array([W])
      else:
          Sigma = D
          W = None
  
      distr = LowrankMultivariateGaussian(
          mu=mx.nd.array([mu]),
          D=mx.nd.array([np.diag(D)]),
          W=W,
          dim=dim,
          rank=rank,
      )
  
      assert np.allclose(
          distr.variance[0].asnumpy(), Sigma, atol=0.1, rtol=0.1
      ), f"did not match: sigma = {Sigma}, sigma_hat = {distr.variance[0]}"
  
      samples = distr.sample(num_samples).squeeze().asnumpy()
  
      theta_hat = maximum_likelihood_estimate_sgd(
          LowrankMultivariateGaussianOutput(
              dim=dim, rank=rank, sigma_init=0.2, sigma_minimum=0.0
          ),
          samples,
          learning_rate=PositiveFloat(0.01),
          num_epochs=PositiveInt(25),
          init_biases=None,  # todo we would need to rework biases a bit to use it in the multivariate case
          hybridize=hybridize,
      )
  
      if rank > 0:
          mu_hat, D_hat, W_hat = theta_hat
          W_hat = mx.nd.array([W_hat])
      else:
          mu_hat, D_hat = theta_hat
          W_hat = None
  
      distr = LowrankMultivariateGaussian(
          dim=dim,
          rank=rank,
          mu=mx.nd.array([mu_hat]),
          D=mx.nd.array([D_hat]),
          W=W_hat,
      )
  
      Sigma_hat = distr.variance.asnumpy()
  
      assert np.allclose(
          mu_hat, mu, atol=0.2, rtol=0.1
      ), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
  
      assert np.allclose(
          Sigma_hat, Sigma, atol=0.1, rtol=0.1
      ), f"sigma did not match: sigma = {Sigma}, sigma_hat = {Sigma_hat}"
  
imports: from functools import reduce
  from typing import List, Tuple
  import mxnet
  import numpy
  import pytest
  from gluonts.mx.model.tpp.distribution import Loglogistic, LoglogisticOutput, Weibull, WeibullOutput
  from gluonts.mx.distribution import Beta, BetaOutput, Binned, BinnedOutput, Categorical, CategoricalOutput, Dirichlet, DirichletMultinomial, DirichletMultinomialOutput, DirichletOutput, DistributionOutput, EmpiricalDistributionOutput, Gamma, GammaOutput, Gaussian, GaussianOutput, GenPareto, GenParetoOutput, Laplace, LaplaceOutput, LogitNormal, LogitNormalOutput, LowrankMultivariateGaussian, LowrankMultivariateGaussianOutput, MultivariateGaussian, MultivariateGaussianOutput, NegativeBinomial, NegativeBinomialOutput, OneInflatedBeta, OneInflatedBetaOutput, PiecewiseLinear, PiecewiseLinearOutput, Poisson, PoissonOutput, StudentT, StudentTOutput, ZeroAndOneInflatedBeta, ZeroAndOneInflatedBetaOutput, ZeroInflatedBeta, ZeroInflatedBetaOutput, ZeroInflatedNegativeBinomialOutput, ZeroInflatedPoissonOutput
  from gluonts.mx.distribution.box_cox_transform import InverseBoxCoxTransform, InverseBoxCoxTransformOutput
  from gluonts.mx.distribution.transformed_distribution import TransformedDistribution
  from gluonts.mx.distribution.transformed_distribution_output import TransformedDistributionOutput
  from gluonts.pydantic import PositiveFloat, PositiveInt
label_desc: Use image node_id 11 for calling a global function with example usage: test_lowrank_multivariate_gaussian(hybridize, rank) without return types