code: def test_image_gradients(batch_size=1, height=5, width=5, channels=1):
      """Test whether the module correctly calculates gradients for known input.
  
      Example input-output pair taken from TF's implementation of image- gradients
  
      """
      image = torch.arange(
          0, batch_size * height * width * channels, dtype=torch.float32
      )
      image = torch.reshape(
          image, (batch_size, channels, height, width)
      )
  
      true_dy = [
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [0.0, 0.0, 0.0, 0.0, 0.0],
      ]
  
      true_dx = [
          [1.0, 1.0, 1.0, 1.0, 0.0],
          [1.0, 1.0, 1.0, 1.0, 0.0],
          [1.0, 1.0, 1.0, 1.0, 0.0],
          [1.0, 1.0, 1.0, 1.0, 0.0],
          [1.0, 1.0, 1.0, 1.0, 0.0],
      ]
  
      true_dy = Tensor(true_dy)
      true_dx = Tensor(true_dx)
  
      dy, dx = image_gradients(image)
  
      assert torch.allclose(dy, true_dy), "dy fails test"
      assert torch.allclose(dx, true_dx), "dx fails tests"
  
imports: import pytest
  import torch
  from torch import Tensor
  from torchmetrics.functional import image_gradients
label_desc: Use image node_id 4 for calling a global function with example usage: test_image_gradients(batch_size, height, width, channels) without return types