code: def test_utilities(head_size):
      size = 5
      cov = torch.randn(size, size)
      cov = torch.mm(cov, cov.t())
  
      mask = torch.ones(size, size)
      mask[head_size:, head_size:] = 0.0
      mask.view(-1)[:: size + 1][head_size:] = 1.0
      arrowhead_full = mask * cov
      expected = torch.flip(
          torch.linalg.cholesky(torch.flip(arrowhead_full, (-2, -1))),
          (-2, -1),
      )
      # test if those flip ops give expected upper triangular values
      assert_close(expected.triu(), expected)
      assert_close(expected.matmul(expected.t()), arrowhead_full)
  
      # test sqrt
      arrowhead = SymmArrowhead(cov[:head_size], cov.diag()[head_size:])
      actual = sqrt(arrowhead)
      assert_close(actual.top, expected[:head_size])
      assert_close(actual.bottom_diag, expected.diag()[head_size:])
  
      # test triu_inverse
      expected = expected.inverse()
      actual = triu_inverse(actual)
      assert_close(actual.top, expected[:head_size])
      assert_close(actual.bottom_diag, expected.diag()[head_size:])
  
      # test triu_matvecmul
      v = torch.randn(size)
      assert_close(triu_matvecmul(actual, v), expected.matmul(v))
      assert_close(
          triu_matvecmul(actual, v, transpose=True),
          expected.t().matmul(v),
      )
  
      # test triu_gram
      actual = triu_gram(actual)
      expected = (
          arrowhead_full.inverse()
          if head_size > 0
          else arrowhead_full.diag().reciprocal()
      )
      assert_close(actual, expected)
  
imports: import pytest
  import torch
  from pyro.ops.arrowhead import SymmArrowhead, sqrt, triu_gram, triu_inverse, triu_matvecmul
  from tests.common import assert_close
label_desc: Use image node_id 1 for calling a global function with example usage: test_utilities(head_size) without return types