code: def test_normal(batch_shape):
      # Test on full data
      data = torch.randn(batch_shape)
      loc = torch.randn(batch_shape).requires_grad_()
      scale = torch.randn(batch_shape).exp().requires_grad_()
      d = dist.NanMaskedNormal(loc, scale)
      d2 = dist.Normal(loc, scale)
      actual = d.log_prob(data)
      expected = d2.log_prob(data)
      assert_close(actual, expected)
  
      # Test on partial data.
      ok = torch.rand(batch_shape) < 0.5
      data[~ok] = math.nan
      actual = d.log_prob(data)
      assert actual.shape == expected.shape
      assert actual.isfinite().all()
      loc_grad, scale_grad = torch.autograd.grad(
          actual.sum(), [loc, scale]
      )
      assert loc_grad.isfinite().all()
      assert scale_grad.isfinite().all()
  
      # Check identity on fully observed and fully unobserved rows.
      assert_close(actual[ok], expected[ok])
      assert_close(actual[~ok], torch.zeros_like(actual[~ok]))
  
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 1 for calling a global function with example usage: test_normal(batch_shape) without return types