text: import warnings
  from typing import cast, Dict, List, Optional, Tuple, Union
  
  import torch
  from torch.nn.functional import pad
  
  from kornia.augmentation.base import GeometricAugmentationBase2D, IntensityAugmentationBase2D, TensorWithTransformMat
  from kornia.color import rgb_to_grayscale
  from kornia.constants import BorderType, pi, Resample, SamplePadding
  from kornia.enhance import (
      adjust_brightness,
      adjust_contrast,
      adjust_hue,
      adjust_saturation,
      equalize,
      invert,
      posterize,
      sharpness,
      solarize,
  )
  from kornia.enhance.normalize import denormalize, normalize
  from kornia.filters import box_blur, gaussian_blur2d, motion_blur
  from kornia.geometry import (
      affine,
      crop_by_transform_mat,
      deg2rad,
      elastic_transform2d,
      get_affine_matrix2d,
      get_perspective_transform,
      get_tps_transform,
      hflip,
      remap,
      resize,
      vflip,
      warp_affine,
      warp_image_tps,
      warp_perspective,
  )
  from kornia.geometry.bbox import bbox_generator, bbox_to_mask
  from kornia.geometry.transform.affwarp import _compute_rotation_matrix, _compute_tensor_center
  from kornia.utils import _extract_device_dtype, create_meshgrid
  
  from . import random_generator as rg
  from .utils import _range_bound, _transform_input
  
  
  class RandomHorizontalFlip(GeometricAugmentationBase2D):
      r"""Apply a random horizontal flip to a tensor image or a batch of tensor images with a given probability.
  
      .. image:: _static/img/RandomHorizontalFlip.png
  
      Input should be a tensor of shape (C, H, W) or a batch of tensors :math:`(B, C, H, W)`.
      If Input is a tuple it is assumed that the first element contains the aforementioned tensors and the second,
      the corresponding transformation matrix that has been applied to them. In this case the module
      will Horizontally flip the tensors and concatenate the corresponding transformation matrix to the
      previous one. This is especially useful when using this functionality as part of an ``nn.Sequential`` module.
  
      Args:
          p: probability of the image being flipped.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.hflip`.
  
      Examples:
          >>> input = torch.tensor([[[[0., 0., 0.],
          ...                         [0., 0., 0.],
          ...                         [0., 1., 1.]]]])
          >>> seq = RandomHorizontalFlip(p=1.0, return_transform=True)
          >>> seq(input)
          (tensor([[[[0., 0., 0.],
                    [0., 0., 0.],
                    [1., 1., 0.]]]]), tensor([[[-1.,  0.,  2.],
                   [ 0.,  1.,  0.],
                   [ 0.,  0.,  1.]]]))
          >>> seq.inverse(seq(input)).equal(input)
          True
      """
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          return dict(batch_shape=torch.as_tensor(batch_shape))
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          w: int = cast(int, params['batch_shape'][-1])
          flip_mat: torch.Tensor = torch.tensor(
              [[-1, 0, w - 1], [0, 1, 0], [0, 0, 1]], device=input.device, dtype=input.dtype
          )
  
          return flip_mat.repeat(input.size(0), 1, 1)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return hflip(input)
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          return self.apply_transform(
              input, params=self._params, transform=torch.as_tensor(transform, device=input.device, dtype=input.dtype)
          )
  
  
  class RandomVerticalFlip(GeometricAugmentationBase2D):
      r"""Apply a random vertical flip to a tensor image or a batch of tensor images with a given probability.
  
      .. image:: _static/img/RandomVerticalFlip.png
  
      Args:
          p: probability of the image being flipped.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                          to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.vflip`.
  
      Examples:
          >>> input = torch.tensor([[[[0., 0., 0.],
          ...                         [0., 0., 0.],
          ...                         [0., 1., 1.]]]])
          >>> seq = RandomVerticalFlip(p=1.0, return_transform=True)
          >>> seq(input)
          (tensor([[[[0., 1., 1.],
                    [0., 0., 0.],
                    [0., 0., 0.]]]]), tensor([[[ 1.,  0.,  0.],
                   [ 0., -1.,  2.],
                   [ 0.,  0.,  1.]]]))
          >>> seq.inverse(seq(input)).equal(input)
          True
      """
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          return dict(batch_shape=torch.as_tensor(batch_shape))
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          h: int = cast(int, params['batch_shape'][-2])
          flip_mat: torch.Tensor = torch.tensor(
              [[1, 0, 0], [0, -1, h - 1], [0, 0, 1]], device=input.device, dtype=input.dtype
          )
  
          return flip_mat.repeat(input.size(0), 1, 1)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return vflip(input)
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          return self.apply_transform(
              input, params=self._params, transform=torch.as_tensor(transform, device=input.device, dtype=input.dtype)
          )
  
  
  class ColorJitter(IntensityAugmentationBase2D):
      r"""Apply a random transformation to the brightness, contrast, saturation and hue of a tensor image.
  
      .. image:: _static/img/ColorJitter.png
  
      Args:
          p: probability of applying the transformation.
          brightness: The brightness factor to apply.
          contrast: The contrast factor to apply.
          saturation: The saturation factor to apply.
          hue: The hue factor to apply.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.enhance.adjust_brightness`,
          :func:`kornia.enhance.adjust_contrast`. :func:`kornia.enhance.adjust_saturation`,
          :func:`kornia.enhance.adjust_hue`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> inputs = torch.ones(1, 3, 3, 3)
          >>> aug = ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.)
          >>> aug(inputs)
          tensor([[[[0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993]],
          <BLANKLINE>
                   [[0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993]],
          <BLANKLINE>
                   [[0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993],
                    [0.9993, 0.9993, 0.9993]]]])
      """
  
      def __init__(
          self,
          brightness: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.0,
          contrast: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.0,
          saturation: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.0,
          hue: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.0,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 1.0,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([brightness, contrast, hue, saturation])
          self.brightness = brightness
          self.contrast = contrast
          self.saturation = saturation
          self.hue = hue
  
      def __repr__(self) -> str:
          repr = f"brightness={self.brightness}, contrast={self.contrast}, saturation={self.saturation}, hue={self.hue}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          brightness: torch.Tensor = _range_bound(
              self.brightness, 'brightness', center=1.0, bounds=(0, 2), device=self._device, dtype=self._dtype
          )
          contrast: torch.Tensor = _range_bound(
              self.contrast, 'contrast', center=1.0, device=self._device, dtype=self._dtype
          )
          saturation: torch.Tensor = _range_bound(
              self.saturation, 'saturation', center=1.0, device=self._device, dtype=self._dtype
          )
          hue: torch.Tensor = _range_bound(self.hue, 'hue', bounds=(-0.5, 0.5), device=self._device, dtype=self._dtype)
          return rg.random_color_jitter_generator(
              batch_shape[0], brightness, contrast, saturation, hue, self.same_on_batch, self.device, self.dtype
          )
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          transforms = [
              lambda img: adjust_brightness(img, params['brightness_factor'] - 1),
              lambda img: adjust_contrast(img, params['contrast_factor']),
              lambda img: adjust_saturation(img, params['saturation_factor']),
              lambda img: adjust_hue(img, params['hue_factor'] * 2 * pi),
          ]
  
          jittered = input
          for idx in params['order'].tolist():
              t = transforms[idx]
              jittered = t(jittered)
  
          return jittered
  
  
  class RandomGrayscale(IntensityAugmentationBase2D):
      r"""Apply random transformation to Grayscale according to a probability p value.
  
      .. image:: _static/img/RandomGrayscale.png
  
      Args:
          p: probability of the image to be transformed to grayscale.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                          to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.color.rgb_to_grayscale`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> inputs = torch.randn((1, 3, 3, 3))
          >>> rec_er = RandomGrayscale(p=1.0)
          >>> rec_er(inputs)
          tensor([[[[-1.1344, -0.1330,  0.1517],
                    [-0.0791,  0.6711, -0.1413],
                    [-0.1717, -0.9023,  0.0819]],
          <BLANKLINE>
                   [[-1.1344, -0.1330,  0.1517],
                    [-0.0791,  0.6711, -0.1413],
                    [-0.1717, -0.9023,  0.0819]],
          <BLANKLINE>
                   [[-1.1344, -0.1330,  0.1517],
                    [-0.0791,  0.6711, -0.1413],
                    [-0.1717, -0.9023,  0.0819]]]])
      """
  
      def __init__(
          self, return_transform: bool = False, same_on_batch: bool = False, p: float = 0.1, keepdim: bool = False
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          # Make sure it returns (*, 3, H, W)
          grayscale = torch.ones_like(input)
          grayscale[:] = rgb_to_grayscale(input)
          return grayscale
  
  
  class RandomErasing(IntensityAugmentationBase2D):
      r"""Erase a random rectangle of a tensor image according to a probability p value.
  
      .. image:: _static/img/RandomErasing.png
  
      The operator removes image parts and fills them with zero values at a selected rectangle
      for each of the images in the batch.
  
      The rectangle will have an area equal to the original image area multiplied by a value uniformly
      sampled between the range [scale[0], scale[1]) and an aspect ratio sampled
      between [ratio[0], ratio[1])
  
      Args:
          p: probability that the random erasing operation will be performed.
          scale: range of proportion of erased area against input image.
          ratio: range of aspect ratio of erased area.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                          to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      Note:
          Input tensor must be float and normalized into [0, 1] for the best differentiability support.
          Additionally, this function accepts another transformation tensor (:math:`(B, 3, 3)`), then the
          applied transformation will be merged int to the input transformation tensor and returned.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> inputs = torch.ones(1, 1, 3, 3)
          >>> rec_er = RandomErasing((.4, .8), (.3, 1/.3), p=0.5)
          >>> rec_er(inputs)
          tensor([[[[1., 0., 0.],
                    [1., 0., 0.],
                    [1., 0., 0.]]]])
      """
  
      # Note: Extra params, inplace=False in Torchvision.
      def __init__(
          self,
          scale: Union[torch.Tensor, Tuple[float, float]] = (0.02, 0.33),
          ratio: Union[torch.Tensor, Tuple[float, float]] = (0.3, 3.3),
          value: float = 0.0,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([scale, ratio])
          self.scale = scale
          self.ratio = ratio
          self.value: float = float(value)
  
      def __repr__(self) -> str:
          repr = f"scale={self.scale}, ratio={self.ratio}, value={self.value}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          scale = torch.as_tensor(self.scale, device=self._device, dtype=self._dtype)
          ratio = torch.as_tensor(self.ratio, device=self._device, dtype=self._dtype)
          return rg.random_rectangles_params_generator(
              batch_shape[0],
              batch_shape[-2],
              batch_shape[-1],
              scale=scale,
              ratio=ratio,
              value=self.value,
              same_on_batch=self.same_on_batch,
              device=self.device,
              dtype=self.dtype,
          )
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          _, c, h, w = input.size()
          values = params['values'].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, *input.shape[1:]).to(input)
  
          bboxes = bbox_generator(params['xs'], params['ys'], params['widths'], params['heights'])
          mask = bbox_to_mask(bboxes, w, h)  # Returns B, H, W
          mask = mask.unsqueeze(1).repeat(1, c, 1, 1).to(input)  # Transform to B, c, H, W
          transformed = torch.where(mask == 1.0, values, input)
          return transformed
  
  
  class RandomPerspective(GeometricAugmentationBase2D):
      r"""Apply a random perspective transformation to an image tensor with a given probability.
  
      .. image:: _static/img/RandomPerspective.png
  
      Args:
          p: probability of the image being perspectively transformed..
          distortion_scale: it controls the degree of distortion and ranges from 0 to 1.
          resample: the interpolation method to use.
          return_transform: if ``True`` return the matrix describing the transformation
                            applied to each.
          same_on_batch: apply the same transformation across the batch. Default: False.
          align_corners: interpolation flag.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.warp_pespective`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> inputs= torch.tensor([[[[1., 0., 0.],
          ...                         [0., 1., 0.],
          ...                         [0., 0., 1.]]]])
          >>> aug = RandomPerspective(0.5, p=0.5)
          >>> out = aug(inputs)
          >>> out
          tensor([[[[0.0000, 0.2289, 0.0000],
                    [0.0000, 0.4800, 0.0000],
                    [0.0000, 0.0000, 0.0000]]]])
          >>> aug.inverse(out)
          tensor([[[[0.0500, 0.0961, 0.0000],
                    [0.2011, 0.3144, 0.0000],
                    [0.0031, 0.0130, 0.0053]]]])
      """
  
      def __init__(
          self,
          distortion_scale: Union[torch.Tensor, float] = 0.5,
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          align_corners: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([distortion_scale])
          self.distortion_scale = distortion_scale
          self.resample: Resample = Resample.get(resample)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
          )
  
      def __repr__(self) -> str:
          repr = (
              f"distortion_scale={self.distortion_scale}, interpolation={self.resample.name}, "
              f"align_corners={self.align_corners}"
          )
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          distortion_scale = torch.as_tensor(self.distortion_scale, device=self._device, dtype=self._dtype)
          return rg.random_perspective_generator(
              batch_shape[0],
              batch_shape[-2],
              batch_shape[-1],
              distortion_scale,
              self.same_on_batch,
              self.device,
              self.dtype,
          )
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return get_perspective_transform(params['start_points'].to(input), params['end_points'].to(input))
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          _, _, height, width = input.shape
          transform = cast(torch.Tensor, transform)
          return warp_perspective(
              input, transform, (height, width), mode=self.resample.name.lower(), align_corners=self.align_corners
          )
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          return self.apply_transform(
              input, params=self._params, transform=torch.as_tensor(transform, device=input.device, dtype=input.dtype)
          )
  
  
  class RandomAffine(GeometricAugmentationBase2D):
      r"""Apply a random 2D affine transformation to a tensor image.
  
      .. image:: _static/img/RandomAffine.png
  
      The transformation is computed so that the image center is kept invariant.
  
      Args:
          p: probability of applying the transformation.
          degrees: Range of degrees to select from.
              If degrees is a number instead of sequence like (min, max), the range of degrees
              will be (-degrees, +degrees). Set to 0 to deactivate rotations.
          translate: tuple of maximum absolute fraction for horizontal
              and vertical translations. For example translate=(a, b), then horizontal shift
              is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
              randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
          scale: scaling factor interval.
              If (a, b) represents isotropic scaling, the scale is randomly sampled from the range a <= scale <= b.
              If (a, b, c, d), the scale is randomly sampled from the range a <= scale_x <= b, c <= scale_y <= d.
              Will keep original scale by default.
          shear: Range of degrees to select from.
              If float, a shear parallel to the x axis in the range (-shear, +shear) will be applied.
              If (a, b), a shear parallel to the x axis in the range (-shear, +shear) will be applied.
              If (a, b, c, d), then x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3])
              will be applied. Will not apply shear by default.
          resample: resample mode from "nearest" (0) or "bilinear" (1).
          padding_mode: padding mode from "zeros" (0), "border" (1) or "refection" (2).
          return_transform: if ``True`` return the matrix describing the transformation applied to each.
          same_on_batch: apply the same transformation across the batch.
          align_corners: interpolation flag.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.warp_affine`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 3, 3)
          >>> aug = RandomAffine((-15., 20.), return_transform=True, p=1.)
          >>> out = aug(input)
          >>> out
          (tensor([[[[0.3961, 0.7310, 0.1574],
                    [0.1781, 0.3074, 0.5648],
                    [0.4804, 0.8379, 0.4234]]]]), tensor([[[ 0.9923, -0.1241,  0.1319],
                   [ 0.1241,  0.9923, -0.1164],
                   [ 0.0000,  0.0000,  1.0000]]]))
          >>> aug.inverse(out)
          tensor([[[[0.3890, 0.6573, 0.1865],
                    [0.2063, 0.3074, 0.5459],
                    [0.3892, 0.7896, 0.4224]]]])
          >>> input
          tensor([[[[0.4963, 0.7682, 0.0885],
                    [0.1320, 0.3074, 0.6341],
                    [0.4901, 0.8964, 0.4556]]]])
      """
  
      def __init__(
          self,
          degrees: Union[torch.Tensor, float, Tuple[float, float]],
          translate: Optional[Union[torch.Tensor, Tuple[float, float]]] = None,
          scale: Optional[Union[torch.Tensor, Tuple[float, float], Tuple[float, float, float, float]]] = None,
          shear: Optional[Union[torch.Tensor, float, Tuple[float, float]]] = None,
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          align_corners: bool = False,
          padding_mode: Union[str, int, SamplePadding] = SamplePadding.ZEROS.name,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([degrees, translate, scale, shear])
          self.degrees = degrees
          self.translate = translate
          self.scale = scale
          self.shear = shear
          self.resample: Resample = Resample.get(resample)
          self.padding_mode: SamplePadding = SamplePadding.get(padding_mode)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              resample=torch.tensor(self.resample.value),
              padding_mode=torch.tensor(self.padding_mode.value),
              align_corners=torch.tensor(align_corners),
          )
  
      def __repr__(self) -> str:
          repr = (
              f"degrees={self.degrees}, translate={self.translate}, scale={self.scale}, shear={self.shear}, "
              f"resample={self.resample.name}"
          )
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          degrees = _range_bound(self.degrees, 'degrees', 0, (-360, 360), device=self._device, dtype=self._dtype)
          translate: Optional[torch.Tensor] = None
          scale: Optional[torch.Tensor] = None
          shear: Optional[torch.Tensor] = None
  
          if self.translate is not None:
              translate = _range_bound(
                  self.translate, 'translate', bounds=(0, 1), check='singular', device=self._device, dtype=self._dtype
              )
          if self.scale is not None:
              scale = torch.as_tensor(self.scale, device=self._device, dtype=self._dtype)
              if len(scale) == 2:
                  scale = _range_bound(
                      scale, 'scale', bounds=(0, float('inf')), check='singular', device=self._device, dtype=self._dtype
                  )
              elif len(scale) == 4:
                  scale = torch.cat(
                      [
                          _range_bound(
                              scale[:2],
                              'scale_x',
                              bounds=(0, float('inf')),
                              check='singular',
                              device=self._device,
                              dtype=self._dtype,
                          ),
                          _range_bound(
                              scale[2:],
                              'scale_y',
                              bounds=(0, float('inf')),
                              check='singular',
                              device=self._device,
                              dtype=self._dtype,
                          ),
                      ]
                  )
              else:
                  raise ValueError(f"'scale' expected to be either 2 or 4 elements. Got {scale}")
          if self.shear is not None:
              shear = torch.as_tensor(self.shear, device=self._device, dtype=self._dtype)
              shear = torch.stack(
                  [
                      _range_bound(
                          shear if shear.dim() == 0 else shear[:2],
                          'shear-x',
                          0,
                          (-360, 360),
                          device=self._device,
                          dtype=self._dtype,
                      ),
                      torch.tensor([0, 0], device=self._device, dtype=self._dtype)
                      if shear.dim() == 0 or len(shear) == 2
                      else _range_bound(shear[2:], 'shear-y', 0, (-360, 360), device=self._device, dtype=self._dtype),
                  ]
              )
          return rg.random_affine_generator(
              batch_shape[0],
              batch_shape[-2],
              batch_shape[-1],
              degrees,
              translate,
              scale,
              shear,
              self.same_on_batch,
              self.device,
              self.dtype,
          )
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return get_affine_matrix2d(
              torch.as_tensor(params['translations'], device=input.device, dtype=input.dtype),
              torch.as_tensor(params['center'], device=input.device, dtype=input.dtype),
              torch.as_tensor(params['scale'], device=input.device, dtype=input.dtype),
              torch.as_tensor(params['angle'], device=input.device, dtype=input.dtype),
              deg2rad(torch.as_tensor(params['sx'], device=input.device, dtype=input.dtype)),
              deg2rad(torch.as_tensor(params['sy'], device=input.device, dtype=input.dtype)),
          )
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          _, _, height, width = input.shape
          transform = cast(torch.Tensor, transform)
          return warp_affine(
              input,
              transform[:, :2, :],
              (height, width),
              self.resample.name.lower(),
              align_corners=self.align_corners,
              padding_mode=self.padding_mode.name.lower(),
          )
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          return self.apply_transform(
              input, params=self._params, transform=torch.as_tensor(transform, device=input.device, dtype=input.dtype)
          )
  
  
  class CenterCrop(GeometricAugmentationBase2D):
      r"""Crop a given image tensor at the center.
  
      .. image:: _static/img/CenterCrop.png
  
      Args:
          size: Desired output size (out_h, out_w) of the crop.
              If integer,  out_h = out_w = size.
              If Tuple[int, int], out_h = size[0], out_w = size[1].
          align_corners: interpolation flag.
          resample: The interpolation mode.
          return_transform: if ``True`` return the matrix describing the transformation
              applied to each.
          p: probability of applying the transformation for the whole batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                          to the batch form (False).
          cropping_mode: The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
                         on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
                         to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
                         differentiability.
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, out_h, out_w)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.crop_by_boxes`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> inputs = torch.randn(1, 1, 4, 4)
          >>> inputs
          tensor([[[[-1.1258, -1.1524, -0.2506, -0.4339],
                    [ 0.8487,  0.6920, -0.3160, -2.1152],
                    [ 0.3223, -1.2633,  0.3500,  0.3081],
                    [ 0.1198,  1.2377,  1.1168, -0.2473]]]])
          >>> aug = CenterCrop(2, p=1., cropping_mode="resample")
          >>> out = aug(inputs)
          >>> out
          tensor([[[[ 0.6920, -0.3160],
                    [-1.2633,  0.3500]]]])
          >>> aug.inverse(out, padding_mode="border")
          tensor([[[[ 0.6920,  0.6920, -0.3160, -0.3160],
                    [ 0.6920,  0.6920, -0.3160, -0.3160],
                    [-1.2633, -1.2633,  0.3500,  0.3500],
                    [-1.2633, -1.2633,  0.3500,  0.3500]]]])
      """
  
      def __init__(
          self,
          size: Union[int, Tuple[int, int]],
          align_corners: bool = True,
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          p: float = 1.0,
          keepdim: bool = False,
          cropping_mode: str = 'slice',
      ) -> None:
          # same_on_batch is always True for CenterCrop
          # Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
          super().__init__(p=1.0, return_transform=return_transform, same_on_batch=True, p_batch=p, keepdim=keepdim)
          if isinstance(size, tuple):
              self.size = (size[0], size[1])
          elif isinstance(size, int):
              self.size = (size, size)
          else:
              raise Exception(f"Invalid size type. Expected (int, tuple(int, int). " f"Got: {type(size)}.")
          self.resample = Resample.get(resample)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
          )
          self.cropping_mode = cropping_mode
  
      def __repr__(self) -> str:
          repr = f"size={self.size}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          return rg.center_crop_generator(batch_shape[0], batch_shape[-2], batch_shape[-1], self.size, self.device)
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          transform: torch.Tensor = get_perspective_transform(params['src'].to(input), params['dst'].to(input))
          transform = transform.expand(input.shape[0], -1, -1)
          return transform
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
              transform = cast(torch.Tensor, transform)
              return crop_by_transform_mat(
                  input, transform[:, :2, :], self.size, self.resample.name.lower(), 'zeros', self.align_corners
              )
          if self.cropping_mode == 'slice':  # uses advanced slicing to crop
              # TODO: implement as separated function `crop_and_resize_iterative`
              B, C, _, _ = input.shape
              H, W = self.size
              out = torch.empty(B, C, H, W, device=input.device, dtype=input.dtype)
              for i in range(B):
                  x1 = int(params['src'][i, 0, 0])
                  x2 = int(params['src'][i, 1, 0]) + 1
                  y1 = int(params['src'][i, 0, 1])
                  y2 = int(params['src'][i, 3, 1]) + 1
                  out[i] = input[i : i + 1, :, y1:y2, x1:x2]
              return out
          raise NotImplementedError(f"Not supported type: {self.cropping_mode}.")
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          if self.cropping_mode != 'resample':
              raise NotImplementedError(
                  f"`inverse` is only applicable for resample cropping mode. Got {self.cropping_mode}."
              )
          if size is None:
              size = self.size
          mode = self.resample.name.lower() if "mode" not in kwargs else kwargs['mode']
          align_corners = self.align_corners if "align_corners" not in kwargs else kwargs['align_corners']
          padding_mode = 'zeros' if "padding_mode" not in kwargs else kwargs['padding_mode']
          transform = cast(torch.Tensor, transform)
          return crop_by_transform_mat(input, transform[:, :2, :], size, mode, padding_mode, align_corners)
  
  
  class RandomRotation(GeometricAugmentationBase2D):
      r"""Apply a random rotation to a tensor image or a batch of tensor images given an amount of degrees.
  
      .. image:: _static/img/RandomRotation.png
  
      Args:
          p: probability of applying the transformation.
          degrees: range of degrees to select from. If degrees is a number the
            range of degrees to select from will be (-degrees, +degrees).
          resample: Default: the interpolation mode.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          align_corners: interpolation flag.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.affine`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.tensor([[1., 0., 0., 2.],
          ...                       [0., 0., 0., 0.],
          ...                       [0., 1., 2., 0.],
          ...                       [0., 0., 1., 2.]])
          >>> aug = RandomRotation(degrees=45.0, return_transform=True, p=1.)
          >>> out = aug(input)
          >>> out[0]
          tensor([[[[0.9824, 0.0088, 0.0000, 1.9649],
                    [0.0000, 0.0029, 0.0000, 0.0176],
                    [0.0029, 1.0000, 1.9883, 0.0000],
                    [0.0000, 0.0088, 1.0117, 1.9649]]]])
          >>> out[1]
          tensor([[[ 1.0000, -0.0059,  0.0088],
                   [ 0.0059,  1.0000, -0.0088],
                   [ 0.0000,  0.0000,  1.0000]]])
          >>> aug.inverse(out)
          tensor([[[[9.6526e-01, 8.6824e-03, 1.7263e-02, 1.9305e+00],
                    [8.6398e-03, 2.9485e-03, 5.8971e-03, 1.7365e-02],
                    [2.9054e-03, 9.9416e-01, 1.9825e+00, 2.3134e-02],
                    [2.5777e-05, 1.1640e-02, 9.9992e-01, 1.9392e+00]]]])
      """
  
      # Note: Extra params, center=None, fill=0 in TorchVision
  
      def __init__(
          self,
          degrees: Union[torch.Tensor, float, Tuple[float, float], List[float]],
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          align_corners: bool = True,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([degrees])
          self.degrees = degrees
          self.resample: Resample = Resample.get(resample)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
          )
  
      def __repr__(self) -> str:
          repr = f"degrees={self.degrees}, interpolation={self.resample.name}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          degrees = _range_bound(self.degrees, 'degrees', 0, (-360, 360), device=self._device, dtype=self._dtype)
          return rg.random_rotation_generator(batch_shape[0], degrees, self.same_on_batch, self.device, self.dtype)
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          # TODO: Update to use `get_rotation_matrix2d`
          angles: torch.Tensor = params["degrees"].to(input)
  
          center: torch.Tensor = _compute_tensor_center(input)
          rotation_mat: torch.Tensor = _compute_rotation_matrix(angles, center.expand(angles.shape[0], -1))
  
          # rotation_mat is B x 2 x 3 and we need a B x 3 x 3 matrix
          trans_mat: torch.Tensor = torch.eye(3, device=input.device, dtype=input.dtype).repeat(input.shape[0], 1, 1)
          trans_mat[:, 0] = rotation_mat[:, 0]
          trans_mat[:, 1] = rotation_mat[:, 1]
  
          return trans_mat
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          transform = cast(torch.Tensor, transform)
          return affine(input, transform[..., :2, :3], self.resample.name.lower(), 'zeros', self.align_corners)
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          return self.apply_transform(
              input, params=self._params, transform=torch.as_tensor(transform, device=input.device, dtype=input.dtype)
          )
  
  
  class RandomCrop(GeometricAugmentationBase2D):
      r"""Crop random patches of a tensor image on a given size.
  
      .. image:: _static/img/RandomCrop.png
  
      Args:
          p: probability of applying the transformation for the whole batch.
          size: Desired output size (out_h, out_w) of the crop.
              Must be Tuple[int, int], then out_h = size[0], out_w = size[1].
          padding: Optional padding on each border
              of the image. Default is None, i.e no padding. If a sequence of length
              4 is provided, it is used to pad left, top, right, bottom borders
              respectively. If a sequence of length 2 is provided, it is used to
              pad left/right, top/bottom borders, respectively.
          pad_if_needed: It will pad the image if smaller than the
              desired size to avoid raising an exception. Since cropping is done
              after padding, the padding seems to be done at a random offset.
          fill: Pixel fill value for constant fill. Default is 0. If a tuple of
              length 3, it is used to fill R, G, B channels respectively.
              This value is only used when the padding_mode is constant.
          padding_mode: Type of padding. Should be: constant, edge, reflect or symmetric.
          resample: the interpolation mode.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          align_corners: interpolation flag.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
          cropping_mode: The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
                         on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
                         to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
                         differentiability.
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, out_h, out_w)`
  
      Note:
          Input tensor must be float and normalized into [0, 1] for the best differentiability support.
          Additionally, this function accepts another transformation tensor (:math:`(B, 3, 3)`), then the
          applied transformation will be merged int to the input transformation tensor and returned.
  
      Examples:
          >>> _ = torch.manual_seed(0)
          >>> inputs = torch.arange(1*1*3*3.).view(1, 1, 3, 3)
          >>> aug = RandomCrop((2, 2), p=1., cropping_mode="resample")
          >>> out = aug(inputs)
          >>> out
          tensor([[[[3., 4.],
                    [6., 7.]]]])
          >>> aug.inverse(out, padding_mode="border")
          tensor([[[[3., 4., 4.],
                    [3., 4., 4.],
                    [6., 7., 7.]]]])
      """
  
      def __init__(
          self,
          size: Tuple[int, int],
          padding: Optional[Union[int, Tuple[int, int], Tuple[int, int, int, int]]] = None,
          pad_if_needed: Optional[bool] = False,
          fill: int = 0,
          padding_mode: str = 'constant',
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          align_corners: bool = True,
          p: float = 1.0,
          keepdim: bool = False,
          cropping_mode: str = 'slice',
      ) -> None:
          # Since PyTorch does not support ragged tensor. So cropping function happens batch-wisely.
          super().__init__(
              p=1.0, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=p, keepdim=keepdim
          )
          self.size = size
          self.padding = padding
          self.pad_if_needed = pad_if_needed
          self.fill = fill
          self.padding_mode = padding_mode
          self.resample: Resample = Resample.get(resample)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
          )
          self.cropping_mode = cropping_mode
  
      def __repr__(self) -> str:
          repr = (
              f"crop_size={self.size}, padding={self.padding}, fill={self.fill}, pad_if_needed={self.pad_if_needed}, "
              f"padding_mode={self.padding_mode}, resample={self.resample.name}"
          )
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          return rg.random_crop_generator(
              batch_shape[0],
              (batch_shape[-2], batch_shape[-1]),
              self.size,
              same_on_batch=self.same_on_batch,
              device=self.device,
              dtype=self.dtype,
          )
  
      def compute_padding(self, shape: torch.Size) -> List[int]:
          if len(shape) != 4:
              raise AssertionError(f"Expected BCHW. Got {shape}.")
          padding = [0, 0, 0, 0]
          if self.padding is not None:
              if isinstance(self.padding, int):
                  self.padding = cast(int, self.padding)
                  padding = [self.padding, self.padding, self.padding, self.padding]
              elif isinstance(self.padding, tuple) and len(self.padding) == 2:
                  self.padding = cast(Tuple[int, int], self.padding)
                  padding = [self.padding[1], self.padding[1], self.padding[0], self.padding[0]]
              elif isinstance(self.padding, tuple) and len(self.padding) == 4:
                  self.padding = cast(Tuple[int, int, int, int], self.padding)
                  padding = [self.padding[3], self.padding[2], self.padding[1], self.padding[0]]
  
          if self.pad_if_needed and shape[-2] < self.size[0]:
              padding = [0, 0, (self.size[0] - shape[-2]), self.size[0] - shape[-2]]
  
          if self.pad_if_needed and shape[-1] < self.size[1]:
              padding = [self.size[1] - shape[-1], self.size[1] - shape[-1], 0, 0]
  
          return padding
  
      def precrop_padding(self, input: torch.Tensor, padding: List[int] = None) -> torch.Tensor:
          if padding is None:
              padding = self.compute_padding(input.shape)
  
          input = pad(input, padding, value=self.fill, mode=self.padding_mode)
  
          return input
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          transform: torch.Tensor = get_perspective_transform(params['src'].to(input), params['dst'].to(input))
          return transform
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
              transform = cast(torch.Tensor, transform)
              return crop_by_transform_mat(
                  input,
                  transform,
                  self.size,
                  mode=self.resample.name.lower(),
                  padding_mode='zeros',
                  align_corners=self.align_corners,
              )
          if self.cropping_mode == 'slice':  # uses advanced slicing to crop
              B, C, _, _ = input.shape
              out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
              for i in range(B):
                  x1 = int(params['src'][i, 0, 0])
                  x2 = int(params['src'][i, 1, 0]) + 1
                  y1 = int(params['src'][i, 0, 1])
                  y2 = int(params['src'][i, 3, 1]) + 1
                  out[i] = input[i : i + 1, :, y1:y2, x1:x2]
              return out
          raise NotImplementedError(f"Not supported type: {self.flags['mode']}.")
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          if self.cropping_mode != 'resample':
              raise NotImplementedError(
                  f"`inverse` is only applicable for resample cropping mode. Got {self.cropping_mode}."
              )
          size = cast(Tuple[int, int], size)
          mode = self.resample.name.lower() if "mode" not in kwargs else kwargs['mode']
          align_corners = self.align_corners if "align_corners" not in kwargs else kwargs['align_corners']
          padding_mode = 'zeros' if "padding_mode" not in kwargs else kwargs['padding_mode']
          transform = cast(torch.Tensor, transform)
          return crop_by_transform_mat(input, transform[:, :2, :], size, mode, padding_mode, align_corners)
  
      def inverse(
          self,
          input: TensorWithTransformMat,
          params: Optional[Dict[str, torch.Tensor]] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          out = super().inverse(input, params, size, **kwargs)
          if params is None:
              params = self._params
          if 'padding_size' in params:
              padding_size = params['padding_size'].unique(dim=0).cpu().squeeze().numpy().tolist()
              padding_size = [-padding_size[0], -padding_size[1], -padding_size[2], -padding_size[3]]
          else:
              padding_size = [0, 0, 0, 0]
          return self.precrop_padding(out, padding_size)
  
      def forward(
          self,
          input: TensorWithTransformMat,
          params: Optional[Dict[str, torch.Tensor]] = None,
          return_transform: Optional[bool] = None,
      ) -> TensorWithTransformMat:
          if isinstance(input, (tuple, list)):
              input_temp = _transform_input(input[0])
              input_pad = self.compute_padding(input[0].shape)
              _input = (self.precrop_padding(input_temp, input_pad), input[1])
          else:
              input = cast(torch.Tensor, input)  # TODO: weird that cast is not working under this context.
              input_temp = _transform_input(input)
              input_pad = self.compute_padding(input_temp.shape)
              _input = self.precrop_padding(input_temp, input_pad)  # type: ignore
          out = super().forward(_input, params, return_transform)
  
          # Update the actual input size for inverse
          _padding_size = torch.tensor(tuple(input_pad), device=input_temp.device, dtype=torch.long).expand(
              input_temp.size(0), -1
          )
          self._params.update({"padding_size": _padding_size})
  
          if not self._params['batch_prob'].all():
              # undo the pre-crop if nothing happened.
              if isinstance(out, tuple) and isinstance(input, tuple):
                  return input[0], out[1]
              if isinstance(out, tuple) and not isinstance(input, tuple):
                  return input, out[1]
              return input
          return out
  
  
  class RandomResizedCrop(GeometricAugmentationBase2D):
      r"""Crop random patches in an image tensor and resizes to a given size.
  
      .. image:: _static/img/RandomResizedCrop.png
  
      Args:
          size: Desired output size (out_h, out_w) of each edge.
              Must be Tuple[int, int], then out_h = size[0], out_w = size[1].
          scale: range of size of the origin size cropped.
          ratio: range of aspect ratio of the origin aspect ratio cropped.
          resample: the interpolation mode.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          align_corners: interpolation flag.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                          to the batch form (False).
          cropping_mode: The used algorithm to crop. ``slice`` will use advanced slicing to extract the tensor based
                         on the sampled indices. ``resample`` will use `warp_affine` using the affine transformation
                         to extract and resize at once. Use `slice` for efficiency, or `resample` for proper
                         differentiability.
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, out_h, out_w)`
  
      Note:
          Input tensor must be float and normalized into [0, 1] for the best differentiability support.
          Additionally, this function accepts another transformation tensor (:math:`(B, 3, 3)`), then the
          applied transformation will be merged int to the input transformation tensor and returned.
  
      Example:
          >>> rng = torch.manual_seed(0)
          >>> inputs = torch.tensor([[[0., 1., 2.],
          ...                         [3., 4., 5.],
          ...                         [6., 7., 8.]]])
          >>> aug = RandomResizedCrop(size=(3, 3), scale=(3., 3.), ratio=(2., 2.), p=1., cropping_mode="resample")
          >>> out = aug(inputs)
          >>> out
          tensor([[[[1.0000, 1.5000, 2.0000],
                    [4.0000, 4.5000, 5.0000],
                    [7.0000, 7.5000, 8.0000]]]])
          >>> aug.inverse(out, padding_mode="border")
          tensor([[[[1., 1., 2.],
                    [4., 4., 5.],
                    [7., 7., 8.]]]])
      """
  
      def __init__(
          self,
          size: Tuple[int, int],
          scale: Union[torch.Tensor, Tuple[float, float]] = (0.08, 1.0),
          ratio: Union[torch.Tensor, Tuple[float, float]] = (3.0 / 4.0, 4.0 / 3.0),
          resample: Union[str, int, Resample] = Resample.BILINEAR.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          align_corners: bool = True,
          p: float = 1.0,
          keepdim: bool = False,
          cropping_mode: str = 'slice',
      ) -> None:
          # Since PyTorch does not support ragged tensor. So cropping function happens all the time.
          super().__init__(
              p=1.0, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=p, keepdim=keepdim
          )
          self._device, self._dtype = _extract_device_dtype([scale, ratio])
          self.size = size
          self.scale = scale
          self.ratio = ratio
          self.resample: Resample = Resample.get(resample)
          self.align_corners = align_corners
          self.flags: Dict[str, torch.Tensor] = dict(
              interpolation=torch.tensor(self.resample.value), align_corners=torch.tensor(align_corners)
          )
          self.cropping_mode = cropping_mode
  
      def __repr__(self) -> str:
          repr = f"size={self.size}, scale={self.scale}, ratio={self.ratio}, interpolation={self.resample.name}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          scale = torch.as_tensor(self.scale, device=self._device, dtype=self._dtype)
          ratio = torch.as_tensor(self.ratio, device=self._device, dtype=self._dtype)
          target_size: torch.Tensor = rg.random_crop_size_generator(
              batch_shape[0],
              (batch_shape[-2], batch_shape[-1]),
              scale,
              ratio,
              self.same_on_batch,
              self.device,
              self.dtype,
          )['size']
          return rg.random_crop_generator(
              batch_shape[0],
              (batch_shape[-2], batch_shape[-1]),
              target_size,
              resize_to=self.size,
              same_on_batch=self.same_on_batch,
              device=self.device,
              dtype=self.dtype,
          )
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          transform: torch.Tensor = get_perspective_transform(params['src'].to(input), params['dst'].to(input))
          transform = transform.expand(input.shape[0], -1, -1)
          return transform
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          if self.cropping_mode == 'resample':  # uses bilinear interpolation to crop
              transform = cast(torch.Tensor, transform)
              return crop_by_transform_mat(
                  input,
                  transform,
                  self.size,
                  mode=self.resample.name.lower(),
                  padding_mode='zeros',
                  align_corners=self.align_corners,
              )
          if self.cropping_mode == 'slice':  # uses advanced slicing to crop
              B, C, _, _ = input.shape
              out = torch.empty(B, C, *self.size, device=input.device, dtype=input.dtype)
              for i in range(B):
                  x1 = int(params['src'][i, 0, 0])
                  x2 = int(params['src'][i, 1, 0]) + 1
                  y1 = int(params['src'][i, 0, 1])
                  y2 = int(params['src'][i, 3, 1]) + 1
                  out[i] = resize(
                      input[i : i + 1, :, y1:y2, x1:x2],
                      self.size,
                      interpolation=(self.resample.name).lower(),
                      align_corners=self.align_corners,
                  )
              return out
          raise NotImplementedError(f"Not supported type: {self.cropping_mode}.")
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs,
      ) -> torch.Tensor:
          if self.cropping_mode != 'resample':
              raise NotImplementedError(
                  f"`inverse` is only applicable for resample cropping mode. Got {self.cropping_mode}."
              )
          size = cast(Tuple[int, int], size)
          mode = self.resample.name.lower() if "mode" not in kwargs else kwargs['mode']
          align_corners = self.align_corners if "align_corners" not in kwargs else kwargs['align_corners']
          padding_mode = 'zeros' if "padding_mode" not in kwargs else kwargs['padding_mode']
          transform = cast(torch.Tensor, transform)
          return crop_by_transform_mat(input, transform[:, :2, :], size, mode, padding_mode, align_corners)
  
  
  class Normalize(IntensityAugmentationBase2D):
      r"""Normalize tensor images with mean and standard deviation.
  
      .. math::
          \text{input[channel] = (input[channel] - mean[channel]) / std[channel]}
  
      Where `mean` is :math:`(M_1, ..., M_n)` and `std` :math:`(S_1, ..., S_n)` for `n` channels,
  
      Args:
          mean: Mean for each channel.
          std: Standard deviations for each channel.
  
      Return:
          Normalised tensor with same size as input :math:`(*, C, H, W)`.
  
      .. note::
          This function internally uses :func:`kornia.enhance.normalize`.
  
      Examples:
  
          >>> norm = Normalize(mean=torch.zeros(4), std=torch.ones(4))
          >>> x = torch.rand(1, 4, 3, 3)
          >>> out = norm(x)
          >>> out.shape
          torch.Size([1, 4, 3, 3])
      """
  
      def __init__(
          self,
          mean: Union[torch.Tensor, Tuple[float], List[float], float],
          std: Union[torch.Tensor, Tuple[float], List[float], float],
          return_transform: bool = False,
          p: float = 1.0,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=True, keepdim=keepdim)
          if isinstance(mean, float):
              mean = torch.tensor([mean])
  
          if isinstance(std, float):
              std = torch.tensor([std])
  
          if isinstance(mean, (tuple, list)):
              mean = torch.tensor(mean)
  
          if isinstance(std, (tuple, list)):
              std = torch.tensor(std)
  
          self.mean = mean
          self.std = std
  
      def __repr__(self) -> str:
          repr = f"mean={self.mean}, std={self.std}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return normalize(input, self.mean, self.std)
  
  
  class Denormalize(IntensityAugmentationBase2D):
      r"""Denormalize tensor images with mean and standard deviation.
  
      .. math::
          \text{input[channel] = (input[channel] * std[channel]) + mean[channel]}
  
      Where `mean` is :math:`(M_1, ..., M_n)` and `std` :math:`(S_1, ..., S_n)` for `n` channels,
  
      Args:
          mean: Mean for each channel.
          std: Standard deviations for each channel.
  
      Return:
          Denormalised tensor with same size as input :math:`(*, C, H, W)`.
  
      .. note::
          This function internally uses :func:`kornia.enhance.denormalize`.
  
      Examples:
  
          >>> norm = Denormalize(mean=torch.zeros(1, 4), std=torch.ones(1, 4))
          >>> x = torch.rand(1, 4, 3, 3)
          >>> out = norm(x)
          >>> out.shape
          torch.Size([1, 4, 3, 3])
      """
  
      def __init__(
          self,
          mean: torch.Tensor,
          std: torch.Tensor,
          return_transform: bool = False,
          p: float = 1.0,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=True, keepdim=keepdim)
          self.mean = mean
          self.std = std
  
      def __repr__(self) -> str:
          repr = f"mean={self.mean}, std={self.std}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return denormalize(input, self.mean, self.std)
  
  
  class RandomMotionBlur(IntensityAugmentationBase2D):
      r"""Perform motion blur on 2D images (4D tensor).
  
      .. image:: _static/img/RandomMotionBlur.png
  
      Args:
          p: probability of applying the transformation.
          kernel_size: motion kernel size (odd and positive).
              If int, the kernel will have a fixed size.
              If Tuple[int, int], it will randomly generate the value from the range batch-wisely.
          angle: angle of the motion blur in degrees (anti-clockwise rotation).
              If float, it will generate the value from (-angle, angle).
          direction: forward/backward direction of the motion blur.
              Lower values towards -1.0 will point the motion blur towards the back (with angle provided via angle),
              while higher values towards 1.0 will point the motion blur forward. A value of 0.0 leads to a
              uniformly (but still angled) motion blur.
              If float, it will generate the value from (-direction, direction).
              If Tuple[int, int], it will randomly generate the value from the range.
          border_type: the padding mode to be applied before convolving.
              CONSTANT = 0, REFLECT = 1, REPLICATE = 2, CIRCULAR = 3.
          resample: the interpolation mode.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      Note:
          Input tensor must be float and normalized into [0, 1] for the best differentiability support.
          Additionally, this function accepts another transformation tensor (:math:`(B, 3, 3)`), then the
          applied transformation will be merged int to the input transformation tensor and returned.
  
          Please set ``resample`` to ``'bilinear'`` if more meaningful gradients wanted.
  
      .. note::
          This function internally uses :func:`kornia.filters.motion_blur`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.ones(1, 1, 5, 5)
          >>> motion_blur = RandomMotionBlur(3, 35., 0.5, p=1.)
          >>> motion_blur(input)
          tensor([[[[0.5773, 1.0000, 1.0000, 1.0000, 0.7561],
                    [0.5773, 1.0000, 1.0000, 1.0000, 0.7561],
                    [0.5773, 1.0000, 1.0000, 1.0000, 0.7561],
                    [0.5773, 1.0000, 1.0000, 1.0000, 0.7561],
                    [0.5773, 1.0000, 1.0000, 1.0000, 0.7561]]]])
      """
  
      def __init__(
          self,
          kernel_size: Union[int, Tuple[int, int]],
          angle: Union[torch.Tensor, float, Tuple[float, float]],
          direction: Union[torch.Tensor, float, Tuple[float, float]],
          border_type: Union[int, str, BorderType] = BorderType.CONSTANT.name,
          resample: Union[str, int, Resample] = Resample.NEAREST.name,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self.kernel_size: Union[int, Tuple[int, int]] = kernel_size
          self._device, self._dtype = _extract_device_dtype([angle, direction])
  
          self.angle = angle
          self.direction = direction
          self.border_type = BorderType.get(border_type)
          self.resample = Resample.get(resample)
          self.flags: Dict[str, torch.Tensor] = {
              "border_type": torch.tensor(self.border_type.value),
              "interpolation": torch.tensor(self.resample.value),
          }
  
      def __repr__(self) -> str:
          repr = (
              f"kernel_size={self.kernel_size}, angle={self.angle}, direction={self.direction}, "
              + f"border_type='{self.border_type.name.lower()}'"
          )
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          angle = _range_bound(
              self.angle, 'angle', center=0.0, bounds=(-360, 360), device=self._device, dtype=self._dtype
          )
          direction = _range_bound(
              self.direction, 'direction', center=0.0, bounds=(-1, 1), device=self._device, dtype=self._dtype
          )
          return rg.random_motion_blur_generator(
              batch_shape[0], self.kernel_size, angle, direction, self.same_on_batch, self.device, self.dtype
          )
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          kernel_size: int = cast(int, params['ksize_factor'].unique().item())
          angle = params['angle_factor']
          direction = params['direction_factor']
          return motion_blur(
              input,
              kernel_size,
              angle,
              direction,
              border_type=self.border_type.name.lower(),
              mode=self.resample.name.lower(),
          )
  
  
  class RandomSolarize(IntensityAugmentationBase2D):
      r"""Solarize given tensor image or a batch of tensor images randomly.
  
      .. image:: _static/img/RandomSolarize.png
  
      Args:
          p: probability of applying the transformation.
          thresholds:
              If float x, threshold will be generated from (0.5 - x, 0.5 + x).
              If tuple (x, y), threshold will be generated from (x, y).
          additions:
              If float x, addition will be generated from (-x, x).
              If tuple (x, y), addition will be generated from (x, y).
          same_on_batch: apply the same transformation across the batch.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.enhance.solarize`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 5, 5)
          >>> solarize = RandomSolarize(0.1, 0.1, p=1.)
          >>> solarize(input)
          tensor([[[[0.4132, 0.1412, 0.1790, 0.2226, 0.3980],
                    [0.2754, 0.4194, 0.0130, 0.4538, 0.2771],
                    [0.4394, 0.4923, 0.1129, 0.2594, 0.3844],
                    [0.3909, 0.2118, 0.1094, 0.2516, 0.3728],
                    [0.2278, 0.0000, 0.4876, 0.0353, 0.5100]]]])
      """
  
      def __init__(
          self,
          thresholds: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.1,
          additions: Union[torch.Tensor, float, Tuple[float, float], List[float]] = 0.1,
          same_on_batch: bool = False,
          return_transform: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch)
          self._device, self._dtype = _extract_device_dtype([thresholds, additions])
          self.thresholds = thresholds
          self.additions = additions
  
      def __repr__(self) -> str:
          repr = f"thresholds={self.thresholds}, additions={self.additions}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          thresholds = _range_bound(
              self.thresholds, 'thresholds', center=0.5, bounds=(0.0, 1.0), device=self._device, dtype=self._dtype
          )
          additions = _range_bound(
              self.additions, 'additions', bounds=(-0.5, 0.5), device=self._device, dtype=self._dtype
          )
          return rg.random_solarize_generator(
              batch_shape[0], thresholds, additions, self.same_on_batch, self.device, self.dtype
          )
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          thresholds = params['thresholds_factor']
          additions: Optional[torch.Tensor]
          if 'additions_factor' in params:
              additions = params['additions_factor']
          else:
              additions = None
          return solarize(input, thresholds, additions)
  
  
  class RandomPosterize(IntensityAugmentationBase2D):
      r"""Posterize given tensor image or a batch of tensor images randomly.
  
      .. image:: _static/img/RandomPosterize.png
  
      Args:
          p: probability of applying the transformation.
          bits: Integer that ranged from (0, 8], in which 0 gives black image and 8 gives the original.
              If int x, bits will be generated from (x, 8).
              If tuple (x, y), bits will be generated from (x, y).
          same_on_batch: apply the same transformation across the batch.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.enhance.posterize`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 5, 5)
          >>> posterize = RandomPosterize(3, p=1.)
          >>> posterize(input)
          tensor([[[[0.4706, 0.7529, 0.0627, 0.1255, 0.2824],
                    [0.6275, 0.4706, 0.8784, 0.4392, 0.6275],
                    [0.3451, 0.3765, 0.0000, 0.1569, 0.2824],
                    [0.5020, 0.6902, 0.7843, 0.1569, 0.2510],
                    [0.6588, 0.9098, 0.3765, 0.8471, 0.4078]]]])
      """
  
      def __init__(
          self,
          bits: Union[int, Tuple[int, int], torch.Tensor] = 3,
          same_on_batch: bool = False,
          return_transform: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([bits])
          self.bits = bits
  
      def __repr__(self) -> str:
          repr = f"(bits={self.bits}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          bits = torch.as_tensor(self.bits, device=self._device, dtype=self._dtype)
          if len(bits.size()) == 0:
              bits = bits.repeat(2)
              bits[1] = 8
          elif not (len(bits.size()) == 1 and bits.size(0) == 2):
              raise ValueError(f"'bits' shall be either a scalar or a length 2 tensor. Got {bits}.")
          return rg.random_posterize_generator(batch_shape[0], bits, self.same_on_batch, self.device, self.dtype)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return posterize(input, params['bits_factor'])
  
  
  class RandomSharpness(IntensityAugmentationBase2D):
      r"""Sharpen given tensor image or a batch of tensor images randomly.
  
      .. image:: _static/img/RandomSharpness.png
  
      Args:
          p: probability of applying the transformation.
          sharpness: factor of sharpness strength. Must be above 0.
          same_on_batch: apply the same transformation across the batch.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.enhance.sharpness`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 5, 5)
          >>> sharpness = RandomSharpness(1., p=1.)
          >>> sharpness(input)
          tensor([[[[0.4963, 0.7682, 0.0885, 0.1320, 0.3074],
                    [0.6341, 0.4810, 0.7367, 0.4177, 0.6323],
                    [0.3489, 0.4428, 0.1562, 0.2443, 0.2939],
                    [0.5185, 0.6462, 0.7050, 0.2288, 0.2823],
                    [0.6816, 0.9152, 0.3971, 0.8742, 0.4194]]]])
      """
  
      def __init__(
          self,
          sharpness: Union[torch.Tensor, float, Tuple[float, float], torch.Tensor] = 0.5,
          same_on_batch: bool = False,
          return_transform: bool = False,
          p: float = 0.5,
          keepdim: bool = False,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
          self._device, self._dtype = _extract_device_dtype([sharpness])
          self.sharpness = sharpness
  
      def __repr__(self) -> str:
          repr = f"sharpness={self.sharpness}"
          return self.__class__.__name__ + f"({repr}, {super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          sharpness = torch.as_tensor(self.sharpness, device=self._device, dtype=self._dtype)
          if sharpness.dim() == 0:
              sharpness = sharpness.repeat(2)
              sharpness[0] = 0.0
          elif not (sharpness.dim() == 1 and sharpness.size(0) == 2):
              raise ValueError(f"'sharpness' must be a scalar or a length 2 tensor. Got {sharpness}.")
          return rg.random_sharpness_generator(batch_shape[0], sharpness, self.same_on_batch, self.device, self.dtype)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          factor = params['sharpness_factor']
          return sharpness(input, factor)
  
  
  class RandomEqualize(IntensityAugmentationBase2D):
      r"""Equalize given tensor image or a batch of tensor images randomly.
  
      .. image:: _static/img/RandomEqualize.png
  
      Args:
          p: Probability to equalize an image.
          same_on_batch: apply the same transformation across the batch.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.enhance.equalize`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 5, 5)
          >>> equalize = RandomEqualize(p=1.)
          >>> equalize(input)
          tensor([[[[0.4963, 0.7682, 0.0885, 0.1320, 0.3074],
                    [0.6341, 0.4901, 0.8964, 0.4556, 0.6323],
                    [0.3489, 0.4017, 0.0223, 0.1689, 0.2939],
                    [0.5185, 0.6977, 0.8000, 0.1610, 0.2823],
                    [0.6816, 0.9152, 0.3971, 0.8742, 0.4194]]]])
      """
  
      def __init__(
          self, same_on_batch: bool = False, return_transform: bool = False, p: float = 0.5, keepdim: bool = False
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return equalize(input)
  
  
  class RandomGaussianBlur(IntensityAugmentationBase2D):
      r"""Apply gaussian blur given tensor image or a batch of tensor images randomly.
  
      .. image:: _static/img/RandomGaussianBlur.png
  
      Args:
          kernel_size: the size of the kernel.
          sigma: the standard deviation of the kernel.
          border_type: the padding mode to be applied before convolving.
            The expected modes are: ``constant``, ``reflect``, ``replicate`` or ``circular``.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`kornia.filters.gaussian_blur2d`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> input = torch.rand(1, 1, 5, 5)
          >>> blur = RandomGaussianBlur((3, 3), (0.1, 2.0), p=1.)
          >>> blur(input)
          tensor([[[[0.6699, 0.4645, 0.3193, 0.1741, 0.1955],
                    [0.5422, 0.6657, 0.6261, 0.6527, 0.5195],
                    [0.3826, 0.2638, 0.1902, 0.1620, 0.2141],
                    [0.6329, 0.6732, 0.5634, 0.4037, 0.2049],
                    [0.8307, 0.6753, 0.7147, 0.5768, 0.7097]]]])
      """
  
      def __init__(
          self,
          kernel_size: Tuple[int, int],
          sigma: Tuple[float, float],
          border_type: str = 'reflect',
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.kernel_size = kernel_size
          self.sigma = sigma
          self.border_type: BorderType = BorderType.get(border_type)
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return gaussian_blur2d(input, self.kernel_size, self.sigma, self.border_type.name.lower())
  
  
  class GaussianBlur(RandomGaussianBlur):
      def __init__(
          self,
          kernel_size: Tuple[int, int],
          sigma: Tuple[float, float],
          border_type: str = 'reflect',
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(
              kernel_size=kernel_size,
              sigma=sigma,
              border_type=border_type,
              return_transform=return_transform,
              same_on_batch=same_on_batch,
              p=p,
          )
          warnings.warn(
              "GaussianBlur is no longer maintained and will be removed from the future versions. "
              "Please use RandomGaussianBlur instead.",
              category=DeprecationWarning,
          )
  
  
  class RandomInvert(IntensityAugmentationBase2D):
      r"""Invert the tensor images values randomly.
  
      .. image:: _static/img/RandomInvert.png
  
      Args:
          max_val: The expected maximum value in the input tensor. The shape has to
            according to the input tensor shape, or at least has to work with broadcasting.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      .. note::
          This function internally uses :func:`kornia.enhance.invert`.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> img = torch.rand(1, 1, 5, 5)
          >>> inv = RandomInvert()
          >>> inv(img)
          tensor([[[[0.4963, 0.7682, 0.0885, 0.1320, 0.3074],
                    [0.6341, 0.4901, 0.8964, 0.4556, 0.6323],
                    [0.3489, 0.4017, 0.0223, 0.1689, 0.2939],
                    [0.5185, 0.6977, 0.8000, 0.1610, 0.2823],
                    [0.6816, 0.9152, 0.3971, 0.8742, 0.4194]]]])
      """
  
      def __init__(
          self,
          max_val: Union[float, torch.Tensor] = torch.tensor(1.0),
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.max_val = max_val
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return invert(input, torch.as_tensor(self.max_val, device=input.device, dtype=input.dtype))
  
  
  class RandomChannelShuffle(IntensityAugmentationBase2D):
      r"""Shuffle the channels of a batch of multi-dimensional images.
  
      .. image:: _static/img/RandomChannelShuffle.png
  
      Args:
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> img = torch.arange(1*2*2*2.).view(1,2,2,2)
          >>> RandomChannelShuffle()(img)
          tensor([[[[4., 5.],
                    [6., 7.]],
          <BLANKLINE>
                   [[0., 1.],
                    [2., 3.]]]])
      """
  
      def __init__(self, return_transform: bool = False, same_on_batch: bool = False, p: float = 0.5) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:
          B, C, _, _ = shape
          channels = torch.rand(B, C).argsort(dim=1)
          return dict(channels=channels)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          out = torch.empty_like(input)
          for i in range(out.shape[0]):
              out[i] = input[i, params['channels'][i]]
          return out
  
  
  class RandomGaussianNoise(IntensityAugmentationBase2D):
      r"""Add gaussian noise to a batch of multi-dimensional images.
  
      .. image:: _static/img/RandomGaussianNoise.png
  
      Args:
          mean: The mean of the gaussian distribution.
          std: The standard deviation of the gaussian distribution.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      Examples:
          >>> rng = torch.manual_seed(0)
          >>> img = torch.ones(1, 1, 2, 2)
          >>> RandomGaussianNoise(mean=0., std=1., p=1.)(img)
          tensor([[[[ 2.5410,  0.7066],
                    [-1.1788,  1.5684]]]])
      """
  
      def __init__(
          self,
          mean: float = 0.0,
          std: float = 1.0,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.mean = mean
          self.std = std
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:
          noise = torch.randn(shape)
          return dict(noise=noise)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return input + params['noise'].to(input.device) * self.std + self.mean
  
  
  class RandomFisheye(GeometricAugmentationBase2D):
      r"""Add random camera radial distortion.
  
      .. image:: _static/img/RandomFisheye.png
  
      Args:
          center_x: Ranges to sample respect to x-coordinate center with shape (2,).
          center_y: Ranges to sample respect to y-coordinate center with shape (2,).
          gamma: Ranges to sample for the gamma values respect to optical center with shape (2,).
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      Examples:
          >>> img = torch.ones(1, 1, 2, 2)
          >>> center_x = torch.tensor([-.3, .3])
          >>> center_y = torch.tensor([-.3, .3])
          >>> gamma = torch.tensor([.9, 1.])
          >>> out = RandomFisheye(center_x, center_y, gamma)(img)
          >>> out.shape
          torch.Size([1, 1, 2, 2])
      """
  
      def __init__(
          self,
          center_x: torch.Tensor,
          center_y: torch.Tensor,
          gamma: torch.Tensor,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.center_x = self._check_tensor(center_x)
          self.center_y = self._check_tensor(center_y)
          self.gamma = self._check_tensor(gamma)
          self.dist = torch.distributions.Uniform
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def _check_tensor(self, data: torch.Tensor) -> torch.Tensor:
          if not isinstance(data, torch.Tensor):
              raise TypeError(f"Invalid input type. Expected torch.Tensor - got: {type(data)}")
  
          if len(data.shape) != 1 and data.shape[0] != 2:
              raise ValueError(f"Tensor must be of shape (2,). Got: {data.shape}.")
  
          return data
  
      def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:
          center_x = self.dist(self.center_x[:1], self.center_x[1:]).rsample(shape[:1])
          center_y = self.dist(self.center_y[:1], self.center_y[1:]).rsample(shape[:1])
          gamma = self.dist(self.gamma[:1], self.gamma[1:]).rsample(shape[:1])
          return dict(center_x=center_x, center_y=center_y, gamma=gamma)
  
      # TODO: It is incorrect to return identity
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return self.identity_matrix(input)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          # create the initial sampling fields
          B, _, H, W = input.shape
          grid = create_meshgrid(H, W, normalized_coordinates=True)
          field_x = grid[..., 0].to(input)  # 1xHxW
          field_y = grid[..., 1].to(input)  # 1xHxW
          # vectorize the random parameters
          center_x = params['center_x'].view(B, 1, 1).to(input)
          center_y = params['center_y'].view(B, 1, 1).to(input)
          gamma = params['gamma'].view(B, 1, 1).to(input)
          # compute and apply the distances respect to the camera optical center
          distance = ((center_x - field_x) ** 2 + (center_y - field_y) ** 2) ** 0.5
          field_x = field_x + field_x * distance ** gamma  # BxHxw
          field_y = field_y + field_y * distance ** gamma  # BxHxW
          return remap(input, field_x, field_y, normalized_coordinates=True, align_corners=True)
  
  
  class RandomElasticTransform(GeometricAugmentationBase2D):
      r"""Add random elastic transformation to a tensor image.
  
      .. image:: _static/img/RandomElasticTransform.png
  
      Args:
          kernel_size: the size of the Gaussian kernel.
          sigma: The standard deviation of the Gaussian in the y and x directions,
            respectively. Larger sigma results in smaller pixel displacements.
          alpha: The scaling factor that controls the intensity of the deformation
            in the y and x directions, respectively.
          align_corners: Interpolation flag used by `grid_sample`.
          mode: Interpolation mode used by `grid_sample`. Either 'bilinear' or 'nearest'.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.elastic_transform2d`.
  
      Examples:
          >>> img = torch.ones(1, 1, 2, 2)
          >>> out = RandomElasticTransform()(img)
          >>> out.shape
          torch.Size([1, 1, 2, 2])
      """
  
      def __init__(
          self,
          kernel_size: Tuple[int, int] = (63, 63),
          sigma: Tuple[float, float] = (32.0, 32.0),
          alpha: Tuple[float, float] = (1.0, 1.0),
          align_corners: bool = False,
          mode: str = 'bilinear',
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.kernel_size = kernel_size
          self.sigma = sigma
          self.alpha = alpha
          self.align_corners = align_corners
          self.mode = mode
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:
          B, _, H, W = shape
          if self.same_on_batch:
              noise = torch.rand(1, 2, H, W, device=self.device, dtype=self.dtype).repeat(B, 1, 1, 1)
          else:
              noise = torch.rand(B, 2, H, W, device=self.device, dtype=self.dtype)
          return dict(noise=noise * 2 - 1)
  
      # TODO: It is incorrect to return identity
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return self.identity_matrix(input)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return elastic_transform2d(
              input, params['noise'].to(input), self.kernel_size, self.sigma, self.alpha, self.align_corners, self.mode
          )
  
  
  class RandomThinPlateSpline(GeometricAugmentationBase2D):
      r"""Add random noise to the Thin Plate Spline algorithm.
  
      .. image:: _static/img/RandomThinPlateSpline.png
  
      Args:
          scale: the scale factor to apply to the destination points.
          align_corners: Interpolation flag used by ``grid_sample``.
          mode: Interpolation mode used by `grid_sample`. Either 'bilinear' or 'nearest'.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      .. note::
          This function internally uses :func:`kornia.geometry.transform.warp_image_tps`.
  
      Examples:
          >>> img = torch.ones(1, 1, 2, 2)
          >>> out = RandomThinPlateSpline()(img)
          >>> out.shape
          torch.Size([1, 1, 2, 2])
      """
  
      def __init__(
          self,
          scale: float = 0.2,
          align_corners: bool = False,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.align_corners = align_corners
          self.dist = torch.distributions.Uniform(-scale, scale)
  
      def generate_parameters(self, shape: torch.Size) -> Dict[str, torch.Tensor]:
          B, _, _, _ = shape
          src = torch.tensor([[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0], [0.0, 0.0]]]).repeat(B, 1, 1)  # Bx5x2
          dst = src + self.dist.rsample(src.shape)
          return dict(src=src, dst=dst)
  
      # TODO: It is incorrect to return identity
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return self.identity_matrix(input)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          src = params['src'].to(input)
          dst = params['dst'].to(input)
          # NOTE: warp_image_tps need to use inverse parameters
          kernel, affine = get_tps_transform(dst, src)
          return warp_image_tps(input, src, kernel, affine, self.align_corners)
  
  
  class RandomBoxBlur(IntensityAugmentationBase2D):
      """Add random blur with a box filter to an image tensor.
  
      .. image:: _static/img/RandomBoxBlur.png
  
      Args:
          kernel_size: the blurring kernel size.
          border_type: the padding mode to be applied before convolving.
            The expected modes are: ``constant``, ``reflect``, ``replicate`` or ``circular``.
          normalized: if True, L1 norm of the kernel is set to 1.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
              input tensor. If ``False`` and the input is a tuple the applied transformation won't be concatenated.
          same_on_batch (bool): apply the same transformation across the batch.
          p: probability of applying the transformation.
  
      .. note::
          This function internally uses :func:`kornia.filters.box_blur`.
  
      Examples:
          >>> img = torch.ones(1, 1, 24, 24)
          >>> out = RandomBoxBlur((7, 7))(img)
          >>> out.shape
          torch.Size([1, 1, 24, 24])
      """
  
      def __init__(
          self,
          kernel_size: Tuple[int, int] = (3, 3),
          border_type: str = 'reflect',
          normalized: bool = True,
          return_transform: bool = False,
          same_on_batch: bool = False,
          p: float = 0.5,
      ) -> None:
          super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, p_batch=1.0)
          self.kernel_size = kernel_size
          self.border_type = border_type
          self.normalized = normalized
  
      def compute_transformation(self, input: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return self.identity_matrix(input)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          return box_blur(input, self.kernel_size, self.border_type, self.normalized)
  
  
  class PadTo(GeometricAugmentationBase2D):
      r"""Pad the given sample to a specific size.
  
      Args:
          size: a tuple of ints in the format (height, width) that give the spatial
              dimensions to pad inputs to.
          pad_mode: the type of padding to perform on the image (valid values
              are those accepted by torch.nn.functional.pad)
          pad_value: fill value for 'constant' padding applied to the image
          p: probability of the image being flipped.
          return_transform: if ``True`` return the matrix describing the transformation applied to each
                            input tensor. If ``False`` and the input is a tuple the applied transformation
                            won't be concatenated.
          same_on_batch: apply the same transformation across the batch.
          keepdim: whether to keep the output shape the same as input (True) or broadcast it
                   to the batch form (False).
  
      Shape:
          - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`, Optional: :math:`(B, 3, 3)`
          - Output: :math:`(B, C, H, W)`
  
      .. note::
          This function internally uses :func:`torch.nn.functional.pad`.
  
      Examples:
          >>> img = torch.tensor([[[[0., 0., 0.],
          ...                       [0., 0., 0.],
          ...                       [0., 0., 0.]]]])
          >>> pad = PadTo((4, 5), pad_value=1.)
          >>> out = pad(img)
          >>> out
          tensor([[[[0., 0., 0., 1., 1.],
                    [0., 0., 0., 1., 1.],
                    [0., 0., 0., 1., 1.],
                    [1., 1., 1., 1., 1.]]]])
          >>> pad.inverse(out)
          tensor([[[[0., 0., 0.],
                    [0., 0., 0.],
                    [0., 0., 0.]]]])
      """
      def __init__(
          self,
          size: Tuple[int, int],
          pad_mode: str = "constant",
          pad_value: Union[int, float] = 0,
          return_transform: bool = False,
      ) -> None:
          super().__init__(p=1., return_transform=return_transform, same_on_batch=True, p_batch=1.)
          self.size = size
          self.pad_mode = pad_mode
          self.pad_value = pad_value
  
      def __repr__(self) -> str:
          return self.__class__.__name__ + f"({super().__repr__()})"
  
      def generate_parameters(self, batch_shape: torch.Size) -> Dict[str, torch.Tensor]:
          input_size = torch.tensor(batch_shape[-2:], dtype=torch.long).expand(batch_shape[0], -1)
          return dict(input_size=input_size)
  
      # TODO: It is incorrect to return identity
      # TODO: Having a resampled version with ``warp_affine``
      def compute_transformation(self, image: torch.Tensor, params: Dict[str, torch.Tensor]) -> torch.Tensor:
          return self.identity_matrix(image)
  
      def apply_transform(
          self, input: torch.Tensor, params: Dict[str, torch.Tensor], transform: Optional[torch.Tensor] = None
      ) -> torch.Tensor:
          _, _, height, width = input.shape
          height_pad: int = self.size[0] - height
          width_pad: int = self.size[1] - width
          return torch.nn.functional.pad(
              input, [0, width_pad, 0, height_pad], mode=self.pad_mode, value=self.pad_value)
  
      def inverse_transform(
          self,
          input: torch.Tensor,
          transform: Optional[torch.Tensor] = None,
          size: Optional[Tuple[int, int]] = None,
          **kwargs
      ) -> torch.Tensor:
          size = cast(Tuple[int, int], size)
          return input[..., :size[0], :size[1]]
  