code: def test_mean_gradient(
      mvn_dist,
      k,
      sample_shape,
      L21,
      omega1,
      L11,
      L22=0.8,
      L33=0.9,
      omega2=0.75,
  ):
      if mvn_dist == "OMTMultivariateNormal" and k > 1:
          return
  
      omega = torch.tensor([omega1, omega2, 0.0])
      loc = torch.zeros(3, requires_grad=True)
      zero_vec = [0.0, 0.0, 0.0]
      off_diag = torch.tensor(
          [zero_vec, [L21, 0.0, 0.0], zero_vec], requires_grad=True
      )
      L = torch.diag(torch.tensor([L11, L22, L33])) + off_diag
  
      if mvn_dist == "OMTMultivariateNormal":
          dist = OMTMultivariateNormal(loc, L)
      elif mvn_dist == "AVFMultivariateNormal":
          CV = (1.1 * torch.rand(2, k, 3)).requires_grad_(True)
          dist = AVFMultivariateNormal(loc, L, CV)
  
      z = dist.rsample(sample_shape)
      torch.cos((omega * z).sum(-1)).mean().backward()
  
      computed_grad = off_diag.grad.cpu().data.numpy()[1, 0]
      analytic = analytic_grad(
          L11=L11, L22=L22, L21=L21, omega1=omega1, omega2=omega2
      )
      assert off_diag.grad.size() == off_diag.size()
      assert loc.grad.size() == loc.size()
      assert torch.triu(off_diag.grad, 1).sum() == 0.0
      assert_equal(
          analytic,
          computed_grad,
          prec=0.005,
          msg="bad cholesky grad for %s (expected %.5f, got %.5f)"
          % (mvn_dist, analytic, computed_grad),
      )
  
file_path: pyro/tests/distributions/test_omt_mvn.py
imports: import numpy
  import pytest
  import torch
  from pyro.distributions import AVFMultivariateNormal, MultivariateNormal, OMTMultivariateNormal
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_mean_gradient(mvn_dist, k, sample_shape, L21, omega1, L11, L22, L33, omega2) without return types