docstr: Args:
      spatial_dims: number of spatial dimensions
      image_size: output spatial size
      decode_size: input spatial size (two or three integers depending on ``spatial_dims``)
      in_channels: number of input channels
      save_theta: whether to save the theta matrix estimation
docstr_tok: ["Args",":","spatial_dims",":","number","of","spatial","dimensions","image_size",":","output","spatial","size","decode_size",":","input","spatial","size","(","two","or","three","integers","depending","on","``","spatial_dims","``",")","in_channels",":","number","of","input","channels","save_theta",":","whether","to","save","the","theta","matrix","estimation"]
code: def __init__(
      self,
      spatial_dims: int,
      image_size: list[int],
      decode_size: list[int],
      in_channels: int,
      save_theta: bool = False,
  ):
      """
      Args:
          spatial_dims: number of spatial dimensions
          image_size: output spatial size
          decode_size: input spatial size (two or three integers depending on ``spatial_dims``)
          in_channels: number of input channels
          save_theta: whether to save the theta matrix estimation
      """
      super().__init__()
      self.spatial_dims = spatial_dims
      if spatial_dims == 2:
          in_features = in_channels * decode_size[0] * decode_size[1]
          out_features = 6
          out_init = torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)
      elif spatial_dims == 3:
          in_features = (
              in_channels
              * decode_size[0]
              * decode_size[1]
              * decode_size[2]
          )
          out_features = 12
          out_init = torch.tensor(
              [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], dtype=torch.float
          )
      else:
          raise ValueError(
              f"only support 2D/3D operation, got spatial_dims={spatial_dims}"
          )
  
      self.fc = nn.Linear(
          in_features=in_features, out_features=out_features
      )
      self.grid = self.get_reference_grid(
          image_size
      )  # (spatial_dims, ...)
  
      # init weight/bias
      self.fc.weight.data.zero_()
      self.fc.bias.data.copy_(out_init)
  
      self.save_theta = save_theta
      self.theta = torch.Tensor()
  
imports: from __future__ import annotations
  import torch
  from torch import nn
  from torch.nn import functional
  from monai.networks.blocks.regunet_block import RegistrationDownSampleBlock, RegistrationExtractionBlock, RegistrationResidualConvBlock, get_conv_block, get_deconv_block
  from monai.networks.utils import meshgrid_ij
label_desc: Use image node_id 1 to create a new AffineHead object from inherited base classes: nn with example: obj = AffineHead(spatial_dims, image_size, decode_size, in_channels, save_theta)