docstr: ``theta`` must be an affine transformation matrix with shape
  3x3 or Nx3x3 or Nx2x3 or 2x3 for spatial 2D transforms,
  4x4 or Nx4x4 or Nx3x4 or 3x4 for spatial 3D transforms,
  where `N` is the batch size. `theta` will be converted into float Tensor for the computation.
  
  Args:
      src (array_like): image in spatial 2D or 3D (N, C, spatial_dims),
          where N is the batch dim, C is the number of channels.
      theta (array_like): Nx3x3, Nx2x3, 3x3, 2x3 for spatial 2D inputs,
          Nx4x4, Nx3x4, 3x4, 4x4 for spatial 3D inputs. When the batch dimension is omitted,
          `theta` will be repeated N times, N is the batch dim of `src`.
      spatial_size: output spatial shape, the full output shape will be
          `[N, C, *spatial_size]` where N and C are inferred from the `src`.
  
  Raises:
      TypeError: When ``theta`` is not a ``torch.Tensor``.
      ValueError: When ``theta`` is not one of [Nxdxd, dxd].
      ValueError: When ``theta`` is not one of [Nx3x3, Nx4x4].
      TypeError: When ``src`` is not a ``torch.Tensor``.
      ValueError: When ``src`` spatially is not one of [2D, 3D].
      ValueError: When affine and image batch dimension differ.
docstr_tok: ["``","theta","``","must","be","an","affine","transformation","matrix","with","shape","3x3","or","Nx3x3","or","Nx2x3","or","2x3","for","spatial","2D","transforms",",","4x4","or","Nx4x4","or","Nx3x4","or","3x4","for","spatial","3D","transforms",",","where","`","N","`","is","the","batch","size",".","`","theta","`","will","be","converted","into","float","Tensor","for","the","computation",".","Args",":","src","(","array_like",")",":","image","in","spatial","2D","or","3D","(","N",",","C",",","spatial_dims",")",",","where","N","is","the","batch","dim",",","C","is","the","number","of","channels",".","theta","(","array_like",")",":","Nx3x3",",","Nx2x3",",","3x3",",","2x3","for","spatial","2D","inputs",",","Nx4x4",",","Nx3x4",",","3x4",",","4x4","for","spatial","3D","inputs",".","When","the","batch","dimension","is","omitted",",","`","theta","`","will","be","repeated","N","times",",","N","is","the","batch","dim","of","`","src","`",".","spatial_size",":","output","spatial","shape",",","the","full","output","shape","will","be","`","[","N",",","C",",","*","spatial_size","]","`","where","N","and","C","are","inferred","from","the","`","src","`",".","Raises",":","TypeError",":","When","``","theta","``","is","not","a","``","torch.Tensor","``",".","ValueError",":","When","``","theta","``","is","not","one","of","[","Nxdxd",",","dxd","]",".","ValueError",":","When","``","theta","``","is","not","one","of","[","Nx3x3",",","Nx4x4","]",".","TypeError",":","When","``","src","``","is","not","a","``","torch.Tensor","``",".","ValueError",":","When","``","src","``","spatially","is","not","one","of","[","2D",",","3D","]",".","ValueError",":","When","affine","and","image","batch","dimension","differ","."]
code: def forward(
      self,
      src: torch.Tensor,
      theta: torch.Tensor,
      spatial_size: Sequence[int] | int | None = None,
  ) -> torch.Tensor:
      """
      ``theta`` must be an affine transformation matrix with shape
      3x3 or Nx3x3 or Nx2x3 or 2x3 for spatial 2D transforms,
      4x4 or Nx4x4 or Nx3x4 or 3x4 for spatial 3D transforms,
      where `N` is the batch size. `theta` will be converted into float Tensor for the computation.
  
      Args:
          src (array_like): image in spatial 2D or 3D (N, C, spatial_dims),
              where N is the batch dim, C is the number of channels.
          theta (array_like): Nx3x3, Nx2x3, 3x3, 2x3 for spatial 2D inputs,
              Nx4x4, Nx3x4, 3x4, 4x4 for spatial 3D inputs. When the batch dimension is omitted,
              `theta` will be repeated N times, N is the batch dim of `src`.
          spatial_size: output spatial shape, the full output shape will be
              `[N, C, *spatial_size]` where N and C are inferred from the `src`.
  
      Raises:
          TypeError: When ``theta`` is not a ``torch.Tensor``.
          ValueError: When ``theta`` is not one of [Nxdxd, dxd].
          ValueError: When ``theta`` is not one of [Nx3x3, Nx4x4].
          TypeError: When ``src`` is not a ``torch.Tensor``.
          ValueError: When ``src`` spatially is not one of [2D, 3D].
          ValueError: When affine and image batch dimension differ.
  
      """
      # validate `theta`
      if not isinstance(theta, torch.Tensor):
          raise TypeError(
              f"theta must be torch.Tensor but is {type(theta).__name__}."
          )
      if theta.dim() not in (2, 3):
          raise ValueError(
              f"theta must be Nxdxd or dxd, got {theta.shape}."
          )
      if theta.dim() == 2:
          theta = theta[None]  # adds a batch dim.
      theta = theta.clone()  # no in-place change of theta
      theta_shape = tuple(theta.shape[1:])
      if theta_shape in ((2, 3), (3, 4)):  # needs padding to dxd
          pad_affine = torch.tensor(
              [0, 0, 1] if theta_shape[0] == 2 else [0, 0, 0, 1]
          )
          pad_affine = pad_affine.repeat(theta.shape[0], 1, 1).to(theta)
          pad_affine.requires_grad = False
          theta = torch.cat([theta, pad_affine], dim=1)
      if tuple(theta.shape[1:]) not in ((3, 3), (4, 4)):
          raise ValueError(
              f"theta must be Nx3x3 or Nx4x4, got {theta.shape}."
          )
      if not torch.is_floating_point(theta):
          raise ValueError(
              f"theta must be floating point data, got {theta.dtype}"
          )
  
      # validate `src`
      if not isinstance(src, torch.Tensor):
          raise TypeError(
              f"src must be torch.Tensor but is {type(src).__name__}."
          )
      sr = src.dim() - 2  # input spatial rank
      if sr not in (2, 3):
          raise ValueError(
              f"Unsupported src dimension: {sr}, available options are [2, 3]."
          )
  
      # set output shape
      src_size = tuple(src.shape)
      dst_size = src_size  # default to the src shape
      if self.spatial_size is not None:
          dst_size = src_size[:2] + self.spatial_size
      if spatial_size is not None:
          dst_size = src_size[:2] + ensure_tuple(spatial_size)
  
      # reverse and normalize theta if needed
      if not self.normalized:
          theta = to_norm_affine(
              affine=theta,
              src_size=src_size[2:],
              dst_size=dst_size[2:],
              align_corners=False,
              zero_centered=self.zero_centered,
          )
      if self.reverse_indexing:
          rev_idx = torch.as_tensor(
              range(sr - 1, -1, -1), device=src.device
          )
          theta[:, :sr] = theta[:, rev_idx]
          theta[:, :, :sr] = theta[:, :, rev_idx]
      if (theta.shape[0] == 1) and src_size[0] > 1:
          # adds a batch dim to `theta` in order to match `src`
          theta = theta.repeat(src_size[0], 1, 1)
      if theta.shape[0] != src_size[0]:
          raise ValueError(
              f"affine and image batch dimension must match, got affine={theta.shape[0]} image={src_size[0]}."
          )
  
      grid = nn.functional.affine_grid(
          theta=theta[:, :sr],
          size=list(dst_size),
          align_corners=self.align_corners,
      )
      dst = nn.functional.grid_sample(
          input=src.contiguous(),
          grid=grid,
          mode=self.mode,
          padding_mode=self.padding_mode,
          align_corners=self.align_corners,
      )
      return dst
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import torch
  import torch.nn
  import monai
  from monai.networks import to_norm_affine
  from monai.utils import GridSampleMode, GridSamplePadMode, convert_to_dst_type, ensure_tuple, look_up_option, optional_import
label_desc: Use image node_id 2 for calling the AffineTransform obj's underlying member method code with example usage: obj.forward(src, theta, spatial_size) and returns: dst