code: def test_multivariate_normal(batch_shape, p):
      # Test on full data
      data = torch.randn(batch_shape + (p,))
      loc = torch.randn(batch_shape + (p,)).requires_grad_()
      scale_tril = torch.randn(batch_shape + (p, p))
      scale_tril.tril_()
      scale_tril.diagonal(dim1=-2, dim2=-1).exp_()
      scale_tril.requires_grad_()
      d = dist.NanMaskedMultivariateNormal(loc, scale_tril=scale_tril)
      d2 = dist.MultivariateNormal(loc, scale_tril=scale_tril)
      actual = d.log_prob(data)
      expected = d2.log_prob(data)
      assert_close(actual, expected)
  
      # Test on partial data.
      ok = torch.rand(batch_shape + (p,)) < 0.5
      data[~ok] = math.nan
      actual = d.log_prob(data)
      assert actual.shape == expected.shape
      assert actual.isfinite().all()
      loc_grad, scale_tril_grad = torch.autograd.grad(
          actual.sum(), [loc, scale_tril]
      )
      assert loc_grad.isfinite().all()
      assert scale_tril_grad.isfinite().all()
  
      # Check identity on fully observed and fully unobserved rows.
      observed = ok.all(-1)
      assert_close(actual[observed], expected[observed])
      unobserved = ~ok.any(-1)
      assert_close(
          actual[unobserved], torch.zeros_like(actual[unobserved])
      )
  
file_path: pyro/tests/distributions/test_nanmasked.py
imports: import math
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.infer import SVI, Trace_ELBO
  from pyro.infer.autoguide import AutoNormal
  from pyro.optim import Adam
  from tests.common import assert_close
label_desc: Use image node_id 2 for calling a global function with example usage: test_multivariate_normal(batch_shape, p) without return types