code: def apply_affine_transform(theta_n, original_grid_homo):
      # theta_n: (N, 2, 3) for 2D, (N, 3, 4) for 3D
      # original_grid_homo: (H, W, 3) for 2D, (D, H, W, 4) for 3D
      assert (
          theta_n.ndim == 3
      ), "theta_n shall have shape of (N, 2, 3) for 2D, (N, 3, 4) for 3D"
      if original_grid_homo.ndim == 3:
          N, dim_2d, dim_homo = theta_n.shape
          assert dim_2d == 2 and dim_homo == 3
          H, W, dim_homo = original_grid_homo.shape
          assert dim_homo == 3
          # reshape to [H * W, dim_homo] and then transpose to [dim_homo, H * W]
          original_grid_transposed = np.transpose(
              np.reshape(original_grid_homo, (H * W, dim_homo))
          )
          grid_n = np.matmul(
              theta_n, original_grid_transposed
          )  # shape (N, dim_2d, H * W)
          # transpose to (N, H * W, dim_2d) and then reshape to (N, H, W, dim_2d)
          grid = np.reshape(
              np.transpose(grid_n, (0, 2, 1)), (N, H, W, dim_2d)
          )
          return grid.astype(np.float32)
      else:
          assert original_grid_homo.ndim == 4
          N, dim_3d, dim_homo = theta_n.shape
          assert dim_3d == 3 and dim_homo == 4
          D, H, W, dim_homo = original_grid_homo.shape
          assert dim_homo == 4
          # reshape to [D * H * W, dim_homo] and then transpose to [dim_homo, D * H * W]
          original_grid_transposed = np.transpose(
              np.reshape(original_grid_homo, (D * H * W, dim_homo))
          )
          grid_n = np.matmul(
              theta_n, original_grid_transposed
          )  # shape (N, dim_3d, D * H * W)
          # transpose to (N, D * H * W, dim_3d) and then reshape to (N, D, H, W, dim_3d)
          grid = np.reshape(
              np.transpose(grid_n, (0, 2, 1)), (N, D, H, W, dim_3d)
          )
          return grid.astype(np.float32)
  
imports: import numpy
  from onnx.reference.op_run import OpRun
label_desc: Use image node_id 2 for calling a global function with example usage: apply_affine_transform(theta_n, original_grid_homo) and returns: grid, grid