Repository: kornia/kornia

Issue:
extract_tensor_patches cuts off parts of the image that would result in a partial patch
### Describe the bug

I first noticed this bug when dividing an image of (1200, 2040) into (256, 256) patches, and then trying to reassemble it. Since the `combine_tensor_patches` function requires the original size, and this did not align with the combined patches, it gave an error:

```
RuntimeError: shape '[-1, 256, 256, 3, 256, 256]' is invalid for input of size 5505024
```

### Reproduction steps

```bash
import torchvision as tv
from PIL import Image
from kornia.contrib import extract_tensor_patches, combine_tensor_patches


crop_size = (256, 256)
image_size = (1200, 2040)
test_image_path = '0804.png'
tensor_transform = tv.transforms.ToTensor()
image = tensor_transform(Image.open(test_image_path))
image = image.unsqueeze(0)
patches = extract_tensor_patches(image, crop_size, stride=256)
reassemble = combine_tensor_patches(patches, original_size=image_size, window_size=crop_size, stride=256)
# see the file used in the next block
```


### Expected behavior

The functions handle patching and reassembly somewhat transparently. 
![0804](https://github.com/kornia/kornia/assets/42846160/d61dbdc6-f18c-4e30-9f61-f76469b85a3c)
The image file I attached is from the DIV2K open source dataset.


### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.27.5
Libc version: glibc-2.35

Python version: 3.11.5 (main, Sep 12 2023, 20:12:53) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.2.0-36-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             8
On-line CPU(s) list:                0-7
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Core(TM) i5-9300H CPU @ 2.40GHz
CPU family:                         6
Model:                              158
Thread(s) per core:                 2
Core(s) per socket:                 4
Socket(s):                          1
Stepping:                           10
CPU max MHz:                        4100.0000
CPU min MHz:                        800.0000
BogoMIPS:                           4800.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp md_clear flush_l1d arch_capabilities
L1d cache:                          128 KiB (4 instances)
L1i cache:                          128 KiB (4 instances)
L2 cache:                           1 MiB (4 instances)
L3 cache:                           8 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-7
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX unsupported
Vulnerability L1tf:                 Mitigation; PTE Inversion
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS, IBPB conditional, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Mitigation; Microcode
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] pytorch-lightning==2.0.9
[pip3] torch==2.0.1
[pip3] torchmetrics==1.1.2
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] Could not collect
```


### Additional context

I'm willing to take a stab at this as a first issue, since I need to fix it for myself anyway. I propose to keep the default behaviour the same, although with a warning, and adding an optional padding to the incomplete patches. 


Discussion:
If you're willing, please open a PR to fix it and we'll interact there, any questions ping us on [Slack](https://join.slack.com/t/kornia/shared_invite/zt-csobk21g-2AQRi~X9Uu6PLMuUZdvfjA)
Worth to check if https://github.com/kornia/kornia/pull/2650 fixes the issue 

Resolving patch:
diff --git a/kornia/contrib/extract_patches.py b/kornia/contrib/extract_patches.py
--- a/kornia/contrib/extract_patches.py
+++ b/kornia/contrib/extract_patches.py
@@ -1,5 +1,6 @@
-import math
+from math import ceil
 from typing import Optional, Tuple, Union, cast
+from warnings import warn
 
 import torch
 import torch.nn.functional as F
@@ -7,19 +8,45 @@
 
 from kornia.core import Module, Tensor, pad
 
+FullPadType = Tuple[int, int, int, int]
+TuplePadType = Union[Tuple[int, int], FullPadType]
+PadType = Union[int, TuplePadType]
+
+
+def create_padding_tuple(padding: PadType, unpadding: bool = False) -> FullPadType:
+    padding = cast(TuplePadType, _pair(padding))
+
+    if len(padding) not in [2, 4]:
+        raise AssertionError(
+            f"{'Unpadding' if unpadding else 'Padding'} must be either an int, tuple of two ints or tuple of four ints"
+        )
+
+    if len(padding) == 2:
+        pad_vert = _pair(padding[0])
+        pad_horz = _pair(padding[1])
+    else:
+        pad_vert = padding[:2]
+        pad_horz = padding[2:]
+    padding = cast(FullPadType, pad_horz + pad_vert)
+
+    return padding
+
 
 def compute_padding(
-    original_size: Union[int, Tuple[int, int]], window_size: Union[int, Tuple[int, int]]
-) -> Tuple[int, int]:
+    original_size: Union[int, Tuple[int, int]],
+    window_size: Union[int, Tuple[int, int]],
+    stride: Optional[Union[int, Tuple[int, int]]] = None,
+) -> FullPadType:
     r"""Compute required padding to ensure chaining of :func:`extract_tensor_patches` and
     :func:`combine_tensor_patches` produces expected result.
 
     Args:
         original_size: the size of the original tensor.
         window_size: the size of the sliding window used while extracting patches.
+        stride: The stride of the sliding window. Optional: if not specified, window_size will be used.
 
     Return:
-        The required padding for `(vertical, horizontal)` as a tuple of 2 ints.
+        The required padding as a tuple of four ints: (top, bottom, left, right)
 
     Example:
         >>> image = torch.arange(12).view(1, 1, 4, 3)
@@ -32,22 +59,44 @@ def compute_padding(
                   [ 9, 10, 11]]]])
 
     .. note::
-        This function is supposed to be used in conjunction with :func:`extract_tensor_patches`
-        and :func:`combine_tensor_patches`.
+        This function will be implicitly used in :func:`extract_tensor_patches` and :func:`combine_tensor_patches` if
+        `allow_auto_(un)padding` is set to True.
     """
     original_size = cast(Tuple[int, int], _pair(original_size))
     window_size = cast(Tuple[int, int], _pair(window_size))
+    if stride is None:
+        stride = window_size
+    stride = cast(Tuple[int, int], _pair(stride))
+
+    remainder_vertical = (original_size[0] - window_size[0]) % stride[0]
+    remainder_horizontal = (original_size[1] - window_size[1]) % stride[1]
+    # it might be best to apply padding only to the far edges (right, bottom), so
+    # that fewer patches are affected by the padding.
+    # For now, just use the default padding
+    if remainder_vertical != 0:
+        vertical_padding = stride[0] - remainder_vertical
+    else:
+        vertical_padding = 0
 
-    if original_size[0] % window_size[0] != 0:
-        padh = math.ceil((window_size[0] - original_size[0] % window_size[0]) / 2)
+    if remainder_horizontal != 0:
+        horizontal_padding = stride[1] - remainder_horizontal
     else:
-        padh = 0
-    if original_size[1] % window_size[1] != 0:
-        padw = math.ceil((window_size[1] - original_size[1] % window_size[1]) / 2)
+        horizontal_padding = 0
+
+    if vertical_padding % 2 == 0:
+        top_padding = bottom_padding = vertical_padding // 2
     else:
-        padw = 0
+        top_padding = vertical_padding // 2
+        bottom_padding = ceil(vertical_padding / 2)
 
-    return (padh, padw)
+    if horizontal_padding % 2 == 0:
+        left_padding = right_padding = horizontal_padding // 2
+    else:
+        left_padding = horizontal_padding // 2
+        right_padding = ceil(horizontal_padding / 2)
+    # the new implementation with unfolding requires symmetric padding
+    padding = int(top_padding), int(bottom_padding), int(left_padding), int(right_padding)
+    return cast(FullPadType, padding)
 
 
 class ExtractTensorPatches(Module):
@@ -71,6 +120,8 @@ class ExtractTensorPatches(Module):
       regulates the overlapping between the extracted patches.
     * :attr:`padding` controls the amount of implicit zeros-paddings on both
       sizes at each dimension.
+    * :attr:`allow_auto_padding` allows automatic calculation of the padding required
+      to fit the window and stride into the image.
 
     The parameters :attr:`window_size`, :attr:`stride` and :attr:`padding` can
     be either:
@@ -80,10 +131,16 @@ class ExtractTensorPatches(Module):
         - a ``tuple`` of two ints -- in which case, the first `int` is used for
           the height dimension, and the second `int` for the width dimension.
 
+    :attr:`padding` can also be a ``tuple`` of four ints -- in which case, the
+    first two ints are for the height dimension while the last two ints are for
+    the width dimension.
+
     Args:
+        input: tensor image where to extract the patches with shape :math:`(B, C, H, W)`.
         window_size: the size of the sliding window and the output patch size.
         stride: stride of the sliding window.
         padding: Zero-padding added to both side of the input.
+        allow_auto_adding: whether to allow automatic padding if the window and stride do not fit into the image.
 
     Shape:
         - Input: :math:`(B, C, H, W)`
@@ -107,20 +164,28 @@ class ExtractTensorPatches(Module):
     def __init__(
         self,
         window_size: Union[int, Tuple[int, int]],
-        stride: Optional[Union[int, Tuple[int, int]]] = 1,
-        padding: Optional[Union[int, Tuple[int, int]]] = 0,
+        stride: Union[int, Tuple[int, int]] = 1,
+        padding: PadType = 0,
+        allow_auto_padding: bool = False,
     ) -> None:
         super().__init__()
-        self.window_size: Tuple[int, int] = _pair(window_size)
-        self.stride: Tuple[int, int] = _pair(stride)
-        self.padding: Tuple[int, int] = _pair(padding)
+        self.window_size: Union[int, Tuple[int, int]] = window_size
+        self.stride: Union[int, Tuple[int, int]] = stride
+        self.padding: PadType = padding
+        self.allow_auto_padding: bool = allow_auto_padding
 
     def forward(self, input: Tensor) -> Tensor:
-        return extract_tensor_patches(input, self.window_size, stride=self.stride, padding=self.padding)
+        return extract_tensor_patches(
+            input,
+            self.window_size,
+            stride=self.stride,
+            padding=self.padding,
+            allow_auto_padding=self.allow_auto_padding,
+        )
 
 
 class CombineTensorPatches(Module):
-    r"""Module that combine patches from tensors.
+    r"""Module that combines patches back into full tensors.
 
     In the simplest case, the output value of the operator with input size
     :math:`(B, N, C, H_{out}, W_{out})` is :math:`(B, C, H, W)`.
@@ -134,14 +199,19 @@ class CombineTensorPatches(Module):
         defined in the function signature.
         left-right and top-bottom order.
 
+
     * :attr:`original_size` is the size of the original image prior to
       extracting tensor patches and defines the shape of the output patch.
     * :attr:`window_size` is the size of the sliding window used while
       extracting tensor patches.
     * :attr:`stride` controls the stride to apply to the sliding window and
       regulates the overlapping between the extracted patches.
-    * :attr:`unpadding` is the amount of padding to be removed. This value
-      must be the same as padding used while extracting tensor patches.
+    * :attr:`unpadding` is the amount of padding to be removed. If specified,
+      this value must be the same as padding used while extracting tensor patches.
+    * :attr:`allow_auto_unpadding` allows automatic calculation of the padding required
+      to fit the window and stride into the image. This must be used if the
+      `allow_auto_padding` flag was used for extracting the patches.
+
 
     The parameters :attr:`original_size`, :attr:`window_size`, :attr:`stride`, and :attr:`unpadding` can
     be either:
@@ -151,12 +221,19 @@ class CombineTensorPatches(Module):
         - a ``tuple`` of two ints -- in which case, the first `int` is used for
           the height dimension, and the second `int` for the width dimension.
 
+    :attr:`unpadding` can also be a ``tuple`` of four ints -- in which case, the
+    first two ints are for the height dimension while the last two ints are for
+    the width dimension.
+
     Args:
-        patches: patched tensor.
-        original_size: the size of the original tensor and the output patch size.
-        window_size: the size of the sliding window used.
+        patches: patched tensor with shape :math:`(B, N, C, H_{out}, W_{out})`.
+        original_size: the size of the original tensor and the output size.
+        window_size: the size of the sliding window used while extracting patches.
         stride: stride of the sliding window.
         unpadding: remove the padding added to both side of the input.
+        allow_auto_unpadding: whether to allow automatic unpadding of the input
+            if the window and stride do not fit into the original_size.
+        eps: small value used to prevent division by zero.
 
     Shape:
         - Input: :math:`(B, N, C, H_{out}, W_{out})`
@@ -176,29 +253,50 @@ class CombineTensorPatches(Module):
 
     def __init__(
         self,
-        original_size: Union[int, Tuple[int, int]],
+        original_size: Tuple[int, int],
         window_size: Union[int, Tuple[int, int]],
-        stride: Union[int, Tuple[int, int]],
-        unpadding: Union[int, Tuple[int, int]] = 0,
+        stride: Optional[Union[int, Tuple[int, int]]] = None,
+        unpadding: PadType = 0,
+        allow_auto_unpadding: bool = False,
     ) -> None:
         super().__init__()
-        self.original_size: Tuple[int, int] = _pair(original_size)
-        self.window_size: Tuple[int, int] = _pair(window_size)
-        self.stride: Tuple[int, int] = _pair(stride)
-        self.unpadding: Tuple[int, int] = _pair(unpadding)
+        self.original_size: Tuple[int, int] = original_size
+        self.window_size: Union[int, Tuple[int, int]] = window_size
+        self.stride: Union[int, Tuple[int, int]] = stride if stride is not None else window_size
+        self.unpadding: PadType = unpadding
+        self.allow_auto_unpadding: bool = allow_auto_unpadding
 
     def forward(self, input: Tensor) -> Tensor:
         return combine_tensor_patches(
-            input, self.original_size, self.window_size, stride=self.stride, unpadding=self.unpadding
+            input,
+            self.original_size,
+            self.window_size,
+            stride=self.stride,
+            unpadding=self.unpadding,
+            allow_auto_unpadding=self.allow_auto_unpadding,
         )
 
 
+def _check_patch_fit(original_size: Tuple[int, int], window_size: Tuple[int, int], stride: Tuple[int, int]) -> bool:
+    remainder_vertical = (original_size[0] - window_size[0]) % stride[0]
+    remainder_horizontal = (original_size[1] - window_size[1]) % stride[1]
+    # the remainder takes into account half a window on each side,
+    # the rest of the image is divided based on the stride, not the window
+    # size
+    if (remainder_horizontal != 0) or (remainder_vertical != 0):
+        # needs padding to fit
+        return False
+    # we can fit a full number of patches in, based on the stride
+    return True
+
+
 def combine_tensor_patches(
     patches: Tensor,
     original_size: Union[int, Tuple[int, int]],
     window_size: Union[int, Tuple[int, int]],
     stride: Union[int, Tuple[int, int]],
-    unpadding: Union[int, Tuple[int, int]] = 0,
+    allow_auto_unpadding: bool = False,
+    unpadding: PadType = 0,
     eps: float = 1e-8,
 ) -> Tensor:
     r"""Restore input from patches.
@@ -207,10 +305,13 @@ def combine_tensor_patches(
 
     Args:
         patches: patched tensor with shape :math:`(B, N, C, H_{out}, W_{out})`.
-        original_size: the size of the original tensor and the output patch size.
+        original_size: the size of the original tensor and the output size.
         window_size: the size of the sliding window used while extracting patches.
         stride: stride of the sliding window.
         unpadding: remove the padding added to both side of the input.
+        allow_auto_unpadding: whether to allow automatic unpadding of the input
+            if the window and stride do not fit into the original_size.
+        eps: small value used to prevent division by zero.
 
     Return:
         The combined patches in an image tensor with shape :math:`(B, C, H, W)`.
@@ -233,10 +334,32 @@ def combine_tensor_patches(
     original_size = cast(Tuple[int, int], _pair(original_size))
     window_size = cast(Tuple[int, int], _pair(window_size))
     stride = cast(Tuple[int, int], _pair(stride))
-    unpadding = cast(Tuple[int, int], _pair(unpadding))
 
-    if len(unpadding) != 2:
-        raise AssertionError("Unpadding must be either an int or a tuple of two ints")
+    if (stride[0] > window_size[0]) | (stride[1] > window_size[1]):
+        raise AssertionError(
+            f"Stride={stride} should be less than or equal to Window size={window_size}, information is missing"
+        )
+
+    if not unpadding:
+        # if padding is specified, we leave it up to the user to ensure it fits
+        # otherwise we check here if it will fit and offer to calculate padding
+        if not _check_patch_fit(original_size, window_size, stride):
+            if not allow_auto_unpadding:
+                warn(
+                    f"The window will not fit into the image. \nWindow size: {window_size}\nStride: {stride}\n"
+                    f"Image size: {original_size}\n"
+                    "This means we probably cannot correctly recombine patches. By enabling `allow_auto_unpadding`, "
+                    "the input will be unpadded to fit the window and stride.\n"
+                    "If the patches have been obtained through `extract_tensor_patches` with the correct padding or "
+                    "the argument `allow_auto_padding`, this will result in a correct reconstruction."
+                )
+            else:
+                unpadding = compute_padding(original_size=original_size, window_size=window_size, stride=stride)
+                # TODO: Can't we just do actual size minus original size to get padding?
+
+    if unpadding:
+        unpadding = create_padding_tuple(unpadding)
+        unpadding = cast(FullPadType, unpadding)
 
     ones = torch.ones(
         patches.shape[0],
@@ -246,6 +369,11 @@ def combine_tensor_patches(
         device=patches.device,
         dtype=patches.dtype,
     )
+
+    if unpadding:
+        ones = pad(ones, pad=unpadding)  # type: ignore
+    restored_size = ones.shape[2:]
+
     patches = patches.permute(0, 2, 3, 4, 1)
     patches = patches.reshape(patches.shape[0], -1, patches.shape[-1])
     int_flag = 0
@@ -256,15 +384,15 @@ def combine_tensor_patches(
         ones = ones.float()
 
     # Calculate normalization map
-    unfold_ones = F.unfold(ones, kernel_size=window_size, stride=stride, padding=unpadding)
-    norm_map = F.fold(
-        input=unfold_ones, output_size=original_size, kernel_size=window_size, stride=stride, padding=unpadding
-    )
+    unfold_ones = F.unfold(ones, kernel_size=window_size, stride=stride)
+    norm_map = F.fold(input=unfold_ones, output_size=restored_size, kernel_size=window_size, stride=stride)
+    if unpadding:
+        norm_map = pad(norm_map, [-i for i in unpadding])  # type: ignore
 
     # Restored tensor
-    saturated_restored_tensor = F.fold(
-        input=patches, output_size=original_size, kernel_size=window_size, stride=stride, padding=unpadding
-    )
+    saturated_restored_tensor = F.fold(input=patches, output_size=restored_size, kernel_size=window_size, stride=stride)
+    if unpadding:
+        saturated_restored_tensor = pad(saturated_restored_tensor, [-i for i in unpadding])  # type: ignore
 
     # Remove satuation effect due to multiple summations
     restored_tensor = saturated_restored_tensor / (norm_map + eps)
@@ -286,9 +414,10 @@ def extract_tensor_patches(
     input: Tensor,
     window_size: Union[int, Tuple[int, int]],
     stride: Union[int, Tuple[int, int]] = 1,
-    padding: Union[int, Tuple[int, int]] = 0,
+    padding: PadType = 0,
+    allow_auto_padding: bool = False,
 ) -> Tensor:
-    r"""Function that extract patches from tensors and stack them.
+    r"""Function that extract patches from tensors and stacks them.
 
     See :class:`~kornia.contrib.ExtractTensorPatches` for details.
 
@@ -297,6 +426,7 @@ def extract_tensor_patches(
         window_size: the size of the sliding window and the output patch size.
         stride: stride of the sliding window.
         padding: Zero-padding added to both side of the input.
+        allow_auto_adding: whether to allow automatic padding if the window and stride do not fit into the image.
 
     Returns:
         the tensor with the extracted patches with shape :math:`(B, N, C, H_{out}, W_{out})`.
@@ -318,21 +448,28 @@ def extract_tensor_patches(
     if len(input.shape) != 4:
         raise ValueError(f"Invalid input shape, we expect BxCxHxW. Got: {input.shape}")
 
+    # check if the window sliding over the image will fit into the image
+    # torch's unfold drops the final patches that don't fit
     window_size = cast(Tuple[int, int], _pair(window_size))
     stride = cast(Tuple[int, int], _pair(stride))
-
-    if (stride[0] > window_size[0]) | (stride[1] > window_size[1]):
-        raise AssertionError(f"Stride={stride} should be less than or equal to Window size={window_size}")
+    original_size = (input.shape[-2], input.shape[-1])
+
+    if not padding:
+        # if padding is specified, we leave it up to the user to ensure it fits
+        # otherwise we check here if it will fit and offer to calculate padding
+        if not _check_patch_fit(original_size, window_size, stride):
+            if not allow_auto_padding:
+                warn(
+                    f"The window will not fit into the image. \nWindow size: {window_size}\nStride: {stride}\n"
+                    f"Image size: {original_size}\n"
+                    "This means that the final incomplete patches will be dropped. By enabling `allow_auto_padding`, "
+                    "the input will be padded to fit the window and stride."
+                )
+            else:
+                padding = compute_padding(original_size=original_size, window_size=window_size, stride=stride)
 
     if padding:
-        padding = cast(Tuple[int, int], _pair(padding))
-
-        if len(padding) != 2:
-            raise AssertionError("Padding must be either an int or a tuple of two ints")
-
-        pad_vert = _pair(padding[0])
-        pad_horz = _pair(padding[1])
-        padding4 = cast(Tuple[int, int, int, int], pad_horz + pad_vert)
-        input = pad(input, padding4)
+        padding = create_padding_tuple(padding)
+        input = pad(input, padding)
 
     return _extract_tensor_patchesnd(input, window_size, stride)


Test patch:
diff --git a/test/contrib/test_contrib.py b/test/contrib/test_contrib.py
--- a/test/contrib/test_contrib.py
+++ b/test/contrib/test_contrib.py
@@ -188,9 +188,9 @@ def test_jit(self, device, dtype):
 
 
 def test_compute_padding():
-    assert kornia.contrib.compute_padding((6, 6), (2, 2)) == (0, 0)
-    assert kornia.contrib.compute_padding((7, 7), (2, 2)) == (1, 1)
-    assert kornia.contrib.compute_padding((8, 7), (4, 4)) == (0, 1)
+    assert kornia.contrib.compute_padding((6, 6), (2, 2)) == (0, 0, 0, 0)
+    assert kornia.contrib.compute_padding((7, 7), (2, 2)) == (0, 1, 0, 1)
+    assert kornia.contrib.compute_padding((8, 7), (4, 4)) == (0, 0, 0, 1)
 
 
 class TestExtractTensorPatches:
@@ -310,6 +310,18 @@ def test_gradcheck(self, device):
         img = utils.tensor_to_gradcheck_var(img)  # to var
         assert gradcheck(kornia.contrib.extract_tensor_patches, (img, 3), raise_exception=True, fast_mode=True)
 
+    def test_auto_padding_stride(self, device, dtype):
+        img_shape = (11, 14)
+        window_size = (3, 3)
+        stride = 2
+        rnge = img_shape[0] * img_shape[1]
+        img = torch.arange(rnge, device=device, dtype=dtype).view(1, 1, *img_shape)
+        patches = kornia.contrib.extract_tensor_patches(
+            img, window_size=window_size, stride=stride, allow_auto_padding=True
+        )
+        # 5 patches vertical, 6 2/3 = 7 horizontal = 35 patches
+        assert patches.shape == (1, 35, 1, *window_size)
+
 
 class TestCombineTensorPatches:
     def test_smoke(self, device, dtype):
@@ -321,10 +333,10 @@ def test_smoke(self, device, dtype):
 
     def test_error(self, device, dtype):
         patches = kornia.contrib.extract_tensor_patches(
-            torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2), padding=1
+            torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(3, 2), stride=(2, 2), padding=1
         )
         with pytest.raises(RuntimeError):
-            kornia.contrib.combine_tensor_patches(patches, original_size=(4, 4), window_size=(2, 2), stride=(3, 2))
+            kornia.contrib.combine_tensor_patches(patches, original_size=(4, 4), window_size=(2, 2), stride=(2, 2))
 
     def test_rect_odd_dim(self, device, dtype):
         img = torch.arange(12, device=device, dtype=dtype).view(1, 1, 4, 3)
@@ -335,15 +347,6 @@ def test_rect_odd_dim(self, device, dtype):
         assert m.shape == (1, 1, 4, 3)
         assert_close(img, m)
 
-    def test_pad_error(self, device, dtype):
-        patches = kornia.contrib.extract_tensor_patches(
-            torch.arange(64, device=device, dtype=dtype).view(1, 1, 8, 8), window_size=(4, 4), stride=(4, 4), padding=1
-        )
-        with pytest.raises(AssertionError):
-            kornia.contrib.combine_tensor_patches(
-                patches, original_size=(8, 8), window_size=(4, 4), stride=(4, 4), unpadding=(1, 1, 1, 1)
-            )
-
     def test_pad_triple_error(self, device, dtype):
         patches = kornia.contrib.extract_tensor_patches(
             torch.arange(36, device=device, dtype=dtype).view(1, 1, 6, 6), window_size=(4, 4), stride=(4, 4), padding=1
@@ -375,22 +378,54 @@ def test_padding2(self, device, dtype):
         assert m(patches).shape == (1, 1, 8, 8)
         assert_close(img, m(patches))
 
+    def test_compute_padding(self, device, dtype):
+        img_shape = (8, 13)
+        rnge = img_shape[0] * img_shape[1]
+        img = torch.arange(rnge, device=device, dtype=dtype).view(1, 1, *img_shape)
+        window_size = (3, 3)
+        padding = kornia.contrib.compute_padding(img_shape, window_size)
+        patches = kornia.contrib.extract_tensor_patches(
+            img, window_size=window_size, stride=window_size, padding=padding
+        )
+        m = kornia.contrib.CombineTensorPatches(img_shape, window_size, stride=window_size, unpadding=padding)
+        assert m(patches).shape == (1, 1, *img_shape)
+        assert_close(img, m(patches))
+
     def test_stride_greater_than_window_size(self, device, dtype):
-        img = torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4)
+        patches = kornia.contrib.extract_tensor_patches(
+            torch.arange(16, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2), padding=1
+        )
         with pytest.raises(AssertionError):
-            kornia.contrib.extract_tensor_patches(img, window_size=(2, 2), stride=(3, 3), padding=1)
+            kornia.contrib.combine_tensor_patches(patches, original_size=(4, 4), window_size=(2, 2), stride=(3, 2))
 
-    def test_autopadding(self, device, dtype):
-        img = torch.arange(104, device=device, dtype=dtype).view(1, 1, 8, 13)
+    def test_expl_autopadding(self, device, dtype):
+        img_shape = (8, 13)
+        img = torch.arange(img_shape[0] * img_shape[1], device=device, dtype=dtype).view(
+            1, 1, img_shape[0], img_shape[1]
+        )
         window_size = (3, 3)
-        padding = kornia.contrib.compute_padding((8, 13), (3, 3))
+        padding = kornia.contrib.compute_padding(img_shape, window_size)
         patches = kornia.contrib.extract_tensor_patches(
             img, window_size=window_size, stride=window_size, padding=padding
         )
-        m = kornia.contrib.CombineTensorPatches((8, 13), (3, 3), stride=window_size, unpadding=padding)
-        assert m(patches).shape == (1, 1, 8, 13)
+        m = kornia.contrib.CombineTensorPatches(img_shape, window_size, stride=window_size, unpadding=padding)
+        assert m(patches).shape == (1, 1, *img_shape)
+
         assert_close(img, m(patches))
 
+    def test_impl_autopadding(self, device, dtype):
+        img_shape = (11, 16)
+        img = torch.arange(img_shape[0] * img_shape[1], device=device, dtype=dtype).view(1, 1, *img_shape)
+        window_size = (3, 3)
+        patches = kornia.contrib.extract_tensor_patches(
+            img, window_size=window_size, stride=window_size, allow_auto_padding=True
+        )
+        recomb = kornia.contrib.combine_tensor_patches(
+            patches, img_shape, window_size=window_size, stride=window_size, allow_auto_unpadding=True
+        )
+        assert recomb.shape == img.shape
+        assert_close(img, recomb)
+
     def test_gradcheck(self, device, dtype):
         patches = kornia.contrib.extract_tensor_patches(
             torch.arange(16.0, device=device, dtype=dtype).view(1, 1, 4, 4), window_size=(2, 2), stride=(2, 2)