code: def test_welford_arrowhead(
      n_samples, dim_size, head_size, regularize
  ):
      adapt_scheme = WelfordArrowheadCovariance(head_size=head_size)
      loc = torch.zeros(dim_size)
      cov = torch.randn(dim_size, dim_size)
      cov = torch.mm(cov, cov.t())
      dist = torch.distributions.MultivariateNormal(
          loc=loc, covariance_matrix=cov
      )
      samples = dist.sample(sample_shape=torch.Size([n_samples]))
  
      for sample in samples:
          adapt_scheme.update(sample)
      top, bottom_diag = adapt_scheme.get_covariance(
          regularize=regularize
      )
      actual = torch.cat(
          [
              top,
              torch.cat(
                  [top[:, head_size:].t(), bottom_diag.diag()], -1
              ),
          ]
      )
  
      mask = torch.ones(dim_size, dim_size)
      mask[head_size:, head_size:] = 0.0
      mask.view(-1)[:: dim_size + 1][head_size:] = 1.0
      expected = np.cov(samples.cpu().numpy(), bias=False, rowvar=False)
      expected = torch.from_numpy(expected).type_as(mask)
      if regularize:
          expected = (
              expected * n_samples + 1e-3 * torch.eye(dim_size) * 5
          ) / (n_samples + 5)
      expected = expected * mask
      assert_equal(actual, expected)
  
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 3 for calling a global function with example usage: test_welford_arrowhead(n_samples, dim_size, head_size, regularize) without return types