code: def check_matrix(self, vals):
      b = zeros((5,))
      for k in range(5):
          b[k] = vals[k, k]
      assert_equal(diag(vals), b)
      b = b * 0
      for k in range(3):
          b[k] = vals[k, k + 2]
      assert_equal(diag(vals, 2), b[:3])
      for k in range(3):
          b[k] = vals[k + 2, k]
      assert_equal(diag(vals, -2), b[:3])
  
imports: import functools
  from unittest import expectedFailure, skipIf
  import pytest
  from pytest import raises
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, parametrize, run_tests, TEST_WITH_TORCHDYNAMO, TestCase, xpassIfTorchDynamo
label_desc: Use image node_id 3 for calling the TestDiag obj's underlying member method code with example usage: obj.check_matrix(vals) without return types