code: def test_welford_diagonal(n_samples, dim_size):
      w = WelfordCovariance(diagonal=True)
      loc = torch.zeros(dim_size)
      cov_diagonal = torch.rand(dim_size)
      cov = torch.diag(cov_diagonal)
      dist = torch.distributions.MultivariateNormal(
          loc=loc, covariance_matrix=cov
      )
      samples = []
      for _ in range(n_samples):
          sample = dist.sample()
          samples.append(sample)
          w.update(sample)
  
      sample_variance = torch.stack(samples).var(dim=0, unbiased=True)
      with optional(pytest.raises(RuntimeError), n_samples == 1):
          estimates = w.get_covariance(regularize=False)
          assert_equal(estimates, sample_variance)
  
file_path: pyro/tests/ops/test_welford.py
imports: import numpy
  import pytest
  import torch
  from pyro.ops.welford import WelfordArrowheadCovariance, WelfordCovariance
  from pyro.util import optional
  from tests.common import assert_equal
label_desc: Use image node_id 1 for calling a global function with example usage: test_welford_diagonal(n_samples, dim_size) without return types