code: def test_shape_extra_dims(self, device, dtype):
      im = torch.rand(1, 1, 3, 5, 5, device=device, dtype=dtype).tile(
          3, 2, 1, 1, 1
      )
      K = torch.rand(1, 1, 3, 3, device=device, dtype=dtype).tile(
          3, 2, 1, 1
      )
      distCoeff = torch.rand(1, 1, 4, device=device, dtype=dtype).tile(
          3, 2, 1
      )
  
      imu = undistort_image(im, K, distCoeff)
      assert imu.shape == (3, 2, 3, 5, 5)
      assert_close(imu[0], imu[1])
  
imports: import pytest
  import torch
  from kornia.geometry.calibration.undistort import undistort_image, undistort_points
  from kornia.testing import assert_close, gradcheck
label_desc: Use image node_id 3 for calling the TestUndistortImage obj's underlying member method code with example usage: obj.test_shape_extra_dims(device, dtype) without return types