code: def test_grad_norm(norm_type, expected):
      """Test utility function for computing the p-norm of individual parameter groups and norm in total."""
  
      class Model(nn.Module):
          def __init__(self):
              super().__init__()
              self.param0 = nn.Parameter(torch.rand(3))
              self.param1 = nn.Parameter(torch.rand(2, 1))
              self.param0.grad = torch.tensor([-1.0, 2.0, -3.0])
              self.param1.grad = torch.tensor([[-4.0], [5.0]])
              # param without grad should not contribute to norm
              self.param2 = nn.Parameter(torch.rand(1))
  
      model = Model()
      norms = grad_norm(model, norm_type)
  
      assert norms.keys() == expected.keys()
      for k in norms:
          assert norms[k] == pytest.approx(expected[k])
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_grads.py
imports: from unittest.mock import Mock
  import pytest
  import torch
  import torch.nn
  from lightning.pytorch.utilities import grad_norm
label_desc: Use image node_id 1 for calling a global function with example usage: test_grad_norm(norm_type, expected) without return types