Repository: kornia/kornia

Issue:
RanAug test failing
### Describe the bug

The test `test_transform_mat` in RandomAugmenter is  failing 
https://github.com/kornia/kornia/blob/main/tests/augmentation/test_auto_operation.py#L119

It seems to be because the operation `get_transformation_matrix` (from parameters) changes the stage of the `aug.transform_matrix`

This version of the test does not fail
![image (51)](https://github.com/kornia/kornia/assets/4803565/2a6e2672-3d7c-425c-bdbd-d2952db13cf8)

Whereas this one, which is current version, does 
![image (50)](https://github.com/kornia/kornia/assets/4803565/1e06eaad-41b8-4277-b6da-901d929abc30)


### Reproduction steps
Just run the tests

```bash
pytest -v --device cpu --dtype float64 tests/augmentation/test_auto_operation.py 
```

And you will see:

```
tests/augmentation/test_auto_operation.py:152: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
testing/base.py:106: in assert_close
    return assert_close(actual, expected, rtol=rtol, atol=atol)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

actual = tensor([[[ 1.0000,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 13.4930],
         [ 0.0000,  0.0000,  1.0000]],

  ...,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 16.0136],
         [ 0.0000,  0.0000,  1.0000]]], dtype=torch.float64)
expected = tensor([[[ 1.0000,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 17.1878],
         [ 0.0000,  0.0000,  1.0000]],

  ...,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 10.7947],
         [ 0.0000,  0.0000,  1.0000]]], dtype=torch.float64)
rtol = 1e-05, atol = 1e-05, kwargs = {}

    def assert_close(
        actual: Tensor, expected: Tensor, *, rtol: Optional[float] = None, atol: Optional[float] = None, **kwargs: Any
    ) -> None:
        if rtol is None and atol is None:
            # `torch.testing.assert_close` used different default tolerances than `torch.testing.assert_allclose`.
            # TODO: remove this special handling as soon as https://github.com/kornia/kornia/issues/1134 is resolved
            #  Basically, this whole wrapper function can be removed and `torch.testing.assert_close` can be used
            #  directly.
            rtol, atol = _default_tolerances(actual, expected)
    
>       return _assert_close(
            actual,
            expected,
            rtol=rtol,
            atol=atol,
            # this is the default value for torch>=1.10, but not for torch==1.9
            # TODO: remove this if kornia relies on torch>=1.10
            check_stride=False,
            equal_nan=False,
            **kwargs,
        )
E       AssertionError: Tensor-likes are not close!
E       
E       Mismatched elements: 10 / 90 (11.1%)
E       Greatest absolute difference: 11.92533016204834 at index (7, 1, 2) (up to 1e-05 allowed)
E       Greatest relative difference: 1.889444158005776 at index (2, 1, 2) (up to 1e-05 allowed)

testing/base.py:37: AssertionError
....

FAILED tests/augmentation/test_auto_operation.py::TestTrivialAugment::test_transform_mat[cpu-float64] - AssertionError: Tensor-likes are not close!
```
```


### Expected behavior

not fail 

### Environment

```shell
Collecting environment information...
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (GCC) 10.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.28

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.19.0-26-cloud-amd64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 12.2.128
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB

Nvidia driver version: 535.86.10
cuDNN version: Could not collect
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
Byte Order:          Little Endian
Address sizes:       46 bits physical, 48 bits virtual
CPU(s):              32
On-line CPU(s) list: 0-31
Thread(s) per core:  2
Core(s) per socket:  16
Socket(s):           1
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               63
Model name:          Intel(R) Xeon(R) CPU @ 2.30GHz
Stepping:            0
CPU MHz:             2299.998
BogoMIPS:            4599.99
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            256K
L3 cache:            46080K
NUMA node0 CPU(s):   0-31
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid xsaveopt arat md_clear arch_capabilities

Versions of relevant libraries:
[pip3] flake8==5.0.4
[pip3] mypy==0.982
[pip3] mypy-extensions==0.4.3
[pip3] natten==0.14.6+torch200cu118
[pip3] numpy==1.24.4
[pip3] numpy-quaternion==2020.9.5.14.42.2
[pip3] onnx==1.14.0
[pip3] onnxruntime==1.14.1
[pip3] open-clip-torch==2.0.2
[pip3] pytest-mypy==0.10.0
[pip3] pytorch-lightning==1.5.0
[pip3] pytorch-metric-learning==2.4.1
[pip3] pytorch-triton==2.1.0+6e4932cda8
[pip3] pytorch3d==0.7.2
[pip3] torch==2.2.1
[pip3] torch-lr-finder==0.2.1
[pip3] torchaudio==2.2.0.dev20231115+cu121
[pip3] torchfile==0.1.0
[pip3] torchgeometry==0.1.2
[pip3] torchmetrics==0.11.4
[pip3] torchsde==0.2.6
[pip3] torchvision==0.17.1
[pip3] triton==2.2.0
[pip3] triton-nightly==2.1.0.dev20230608213938
[conda] natten                    0.14.6+torch200cu118          pypi_0    pypi
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpy-quaternion          2020.9.5.14.42.2          pypi_0    pypi
[conda] open-clip-torch           2.0.2                    pypi_0    pypi
[conda] pytorch-lightning         1.5.0                    pypi_0    pypi
[conda] pytorch-metric-learning   2.4.1                    pypi_0    pypi
[conda] pytorch-triton            2.1.0+6e4932cda8          pypi_0    pypi
[conda] pytorch3d                 0.7.2                    pypi_0    pypi
[conda] torch                     2.2.1                    pypi_0    pypi
[conda] torch-lr-finder           0.2.1                    pypi_0    pypi
[conda] torchaudio                2.2.0.dev20231115+cu121          pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchgeometry             0.1.2                    pypi_0    pypi
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchsde                  0.2.6                    pypi_0    pypi
[conda] torchvision               0.17.1                   pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi
[conda] triton-nightly            2.1.0.dev20230608213938          pypi_0    pypi
```
```


### Additional context

_No response_
RanAug test failing
### Describe the bug

The test `test_transform_mat` in RandomAugmenter is  failing 
https://github.com/kornia/kornia/blob/main/tests/augmentation/test_auto_operation.py#L119

It seems to be because the operation `get_transformation_matrix` (from parameters) changes the stage of the `aug.transform_matrix`

This version of the test does not fail
![image (51)](https://github.com/kornia/kornia/assets/4803565/2a6e2672-3d7c-425c-bdbd-d2952db13cf8)

Whereas this one, which is current version, does 
![image (50)](https://github.com/kornia/kornia/assets/4803565/1e06eaad-41b8-4277-b6da-901d929abc30)


### Reproduction steps
Just run the tests

```bash
pytest -v --device cpu --dtype float64 tests/augmentation/test_auto_operation.py 
```

And you will see:

```
tests/augmentation/test_auto_operation.py:152: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
testing/base.py:106: in assert_close
    return assert_close(actual, expected, rtol=rtol, atol=atol)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

actual = tensor([[[ 1.0000,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 13.4930],
         [ 0.0000,  0.0000,  1.0000]],

  ...,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 16.0136],
         [ 0.0000,  0.0000,  1.0000]]], dtype=torch.float64)
expected = tensor([[[ 1.0000,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 17.1878],
         [ 0.0000,  0.0000,  1.0000]],

  ...,  0.0000,  0.0000],
         [ 0.0000,  1.0000, 10.7947],
         [ 0.0000,  0.0000,  1.0000]]], dtype=torch.float64)
rtol = 1e-05, atol = 1e-05, kwargs = {}

    def assert_close(
        actual: Tensor, expected: Tensor, *, rtol: Optional[float] = None, atol: Optional[float] = None, **kwargs: Any
    ) -> None:
        if rtol is None and atol is None:
            # `torch.testing.assert_close` used different default tolerances than `torch.testing.assert_allclose`.
            # TODO: remove this special handling as soon as https://github.com/kornia/kornia/issues/1134 is resolved
            #  Basically, this whole wrapper function can be removed and `torch.testing.assert_close` can be used
            #  directly.
            rtol, atol = _default_tolerances(actual, expected)
    
>       return _assert_close(
            actual,
            expected,
            rtol=rtol,
            atol=atol,
            # this is the default value for torch>=1.10, but not for torch==1.9
            # TODO: remove this if kornia relies on torch>=1.10
            check_stride=False,
            equal_nan=False,
            **kwargs,
        )
E       AssertionError: Tensor-likes are not close!
E       
E       Mismatched elements: 10 / 90 (11.1%)
E       Greatest absolute difference: 11.92533016204834 at index (7, 1, 2) (up to 1e-05 allowed)
E       Greatest relative difference: 1.889444158005776 at index (2, 1, 2) (up to 1e-05 allowed)

testing/base.py:37: AssertionError
....

FAILED tests/augmentation/test_auto_operation.py::TestTrivialAugment::test_transform_mat[cpu-float64] - AssertionError: Tensor-likes are not close!
```
```


### Expected behavior

not fail 

### Environment

```shell
Collecting environment information...
PyTorch version: 2.2.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (GCC) 10.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.28

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.19.0-26-cloud-amd64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 12.2.128
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB

Nvidia driver version: 535.86.10
cuDNN version: Could not collect
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
Byte Order:          Little Endian
Address sizes:       46 bits physical, 48 bits virtual
CPU(s):              32
On-line CPU(s) list: 0-31
Thread(s) per core:  2
Core(s) per socket:  16
Socket(s):           1
NUMA node(s):        1
Vendor ID:           GenuineIntel
CPU family:          6
Model:               63
Model name:          Intel(R) Xeon(R) CPU @ 2.30GHz
Stepping:            0
CPU MHz:             2299.998
BogoMIPS:            4599.99
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            256K
L3 cache:            46080K
NUMA node0 CPU(s):   0-31
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid xsaveopt arat md_clear arch_capabilities

Versions of relevant libraries:
[pip3] flake8==5.0.4
[pip3] mypy==0.982
[pip3] mypy-extensions==0.4.3
[pip3] natten==0.14.6+torch200cu118
[pip3] numpy==1.24.4
[pip3] numpy-quaternion==2020.9.5.14.42.2
[pip3] onnx==1.14.0
[pip3] onnxruntime==1.14.1
[pip3] open-clip-torch==2.0.2
[pip3] pytest-mypy==0.10.0
[pip3] pytorch-lightning==1.5.0
[pip3] pytorch-metric-learning==2.4.1
[pip3] pytorch-triton==2.1.0+6e4932cda8
[pip3] pytorch3d==0.7.2
[pip3] torch==2.2.1
[pip3] torch-lr-finder==0.2.1
[pip3] torchaudio==2.2.0.dev20231115+cu121
[pip3] torchfile==0.1.0
[pip3] torchgeometry==0.1.2
[pip3] torchmetrics==0.11.4
[pip3] torchsde==0.2.6
[pip3] torchvision==0.17.1
[pip3] triton==2.2.0
[pip3] triton-nightly==2.1.0.dev20230608213938
[conda] natten                    0.14.6+torch200cu118          pypi_0    pypi
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpy-quaternion          2020.9.5.14.42.2          pypi_0    pypi
[conda] open-clip-torch           2.0.2                    pypi_0    pypi
[conda] pytorch-lightning         1.5.0                    pypi_0    pypi
[conda] pytorch-metric-learning   2.4.1                    pypi_0    pypi
[conda] pytorch-triton            2.1.0+6e4932cda8          pypi_0    pypi
[conda] pytorch3d                 0.7.2                    pypi_0    pypi
[conda] torch                     2.2.1                    pypi_0    pypi
[conda] torch-lr-finder           0.2.1                    pypi_0    pypi
[conda] torchaudio                2.2.0.dev20231115+cu121          pypi_0    pypi
[conda] torchfile                 0.1.0                    pypi_0    pypi
[conda] torchgeometry             0.1.2                    pypi_0    pypi
[conda] torchmetrics              0.11.4                   pypi_0    pypi
[conda] torchsde                  0.2.6                    pypi_0    pypi
[conda] torchvision               0.17.1                   pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi
[conda] triton-nightly            2.1.0.dev20230608213938          pypi_0    pypi
```
```


### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/kornia/augmentation/auto/base.py b/kornia/augmentation/auto/base.py
--- a/kornia/augmentation/auto/base.py
+++ b/kornia/augmentation/auto/base.py
@@ -15,7 +15,7 @@
 SUBPOLICY_CONFIG = List[OP_CONFIG]
 
 
-class PolicyAugmentBase(TransformMatrixMinIn, ImageSequentialBase):
+class PolicyAugmentBase(ImageSequentialBase, TransformMatrixMinIn):
     """Policy-based image augmentation."""
 
     def __init__(self, policy: List[SUBPOLICY_CONFIG], transformation_matrix_mode: str = "silence") -> None:
@@ -92,8 +92,25 @@ def forward_parameters(self, batch_shape: torch.Size) -> List[ParamItem]:
         return params
 
     def transform_inputs(self, input: Tensor, params: List[ParamItem], extra_args: Dict[str, Any] = {}) -> Tensor:
+        for param in params:
+            module = self.get_submodule(param.name)
+            input = InputSequentialOps.transform(input, module=module, param=param, extra_args=extra_args)
+        return input
+
+    def forward(
+        self, input: Tensor, params: Optional[List[ParamItem]] = None, extra_args: Dict[str, Any] = {}
+    ) -> Tensor:
+        self.clear_state()
+
+        if params is None:
+            inp = input
+            _, out_shape = self.autofill_dim(inp, dim_range=(2, 4))
+            params = self.forward_parameters(out_shape)
+
         for param in params:
             module = self.get_submodule(param.name)
             input = InputSequentialOps.transform(input, module=module, param=param, extra_args=extra_args)
             self._update_transform_matrix_by_module(module)
+
+        self._params = params
         return input
diff --git a/kornia/augmentation/auto/operations/policy.py b/kornia/augmentation/auto/operations/policy.py
--- a/kornia/augmentation/auto/operations/policy.py
+++ b/kornia/augmentation/auto/operations/policy.py
@@ -74,9 +74,9 @@ def get_transformation_matrix(
                     flags = override_parameters(module.op.flags, extra_args, in_place=False)
                     mat = module.op.generate_transformation_matrix(input, param.data, flags)
                 elif module.op._transform_matrix is not None:
-                    mat = as_tensor(module.op._transform_matrix, device=input.device, dtype=input.dtype)
+                    mat = as_tensor(module.transform_matrix, device=input.device, dtype=input.dtype)
                 else:
-                    raise RuntimeError(f"{module}.op._transform_matrix is None while `recompute=False`.")
+                    raise RuntimeError(f"{module}.transform_matrix is None while `recompute=False`.")
                 res_mat = mat @ res_mat
                 input = module.op.transform_output_tensor(input, ori_shape)
                 if module.op.keepdim and ori_shape != input.shape:
diff --git a/kornia/augmentation/auto/rand_augment/rand_augment.py b/kornia/augmentation/auto/rand_augment/rand_augment.py
--- a/kornia/augmentation/auto/rand_augment/rand_augment.py
+++ b/kornia/augmentation/auto/rand_augment/rand_augment.py
@@ -1,7 +1,6 @@
 from typing import Dict, Iterator, List, Optional, Tuple, Union, cast
 
 import torch
-from torch.distributions import Categorical
 
 from kornia.augmentation.auto.base import SUBPOLICY_CONFIG, PolicyAugmentBase
 from kornia.augmentation.auto.operations import OperationBase
@@ -70,11 +69,14 @@ def __init__(
             _policy = policy
 
         super().__init__(_policy, transformation_matrix_mode=transformation_matrix_mode)
-        selection_weights = torch.tensor([1.0 / len(self)] * len(self))
-        self.rand_selector = Categorical(selection_weights)
         self.n = n
         self.m = m
 
+    def rand_selector(self, n: int) -> Tensor:
+        perm = torch.randperm(len(self._modules))
+        idx = perm[:n]
+        return idx
+
     def compose_subpolicy_sequential(self, subpolicy: SUBPOLICY_CONFIG) -> PolicySequential:
         if len(subpolicy) != 1:
             raise RuntimeError(f"Each policy must have only one operation for RandAugment. Got {len(subpolicy)}.")
@@ -83,14 +85,15 @@ def compose_subpolicy_sequential(self, subpolicy: SUBPOLICY_CONFIG) -> PolicySeq
 
     def get_forward_sequence(self, params: Optional[List[ParamItem]] = None) -> Iterator[Tuple[str, Module]]:
         if params is None:
-            idx = self.rand_selector.sample((self.n,))
+            idx = self.rand_selector(
+                self.n,
+            )
             return self.get_children_by_indices(idx)
 
         return self.get_children_by_params(params)
 
     def forward_parameters(self, batch_shape: torch.Size) -> List[ParamItem]:
         named_modules: Iterator[Tuple[str, Module]] = self.get_forward_sequence()
-
         params: List[ParamItem] = []
         mod_param: Union[Dict[str, Tensor], List[ParamItem]]
         m = torch.tensor([self.m / 30] * batch_shape[0])


Test patch:
diff --git a/tests/augmentation/test_auto_operation.py b/tests/augmentation/test_auto_operation.py
--- a/tests/augmentation/test_auto_operation.py
+++ b/tests/augmentation/test_auto_operation.py
@@ -11,7 +11,6 @@
 from kornia.augmentation.auto.trivial_augment import TrivialAugment
 from kornia.augmentation.container import AugmentationSequential
 from kornia.geometry.bbox import bbox_to_mask
-from kornia.utils._compat import torch_version
 
 from testing.base import BaseTester
 
@@ -112,10 +111,6 @@ def test_smoke(self, policy):
         in_tensor = torch.rand(10, 3, 50, 50, requires_grad=True)
         aug(in_tensor)
 
-    @pytest.mark.xfail(
-        torch_version() in {"1.9.1", "1.10.2", "1.11.0", "1.12.1", "1.13.1"},
-        reason="randomness failing into some torch versions",
-    )
     def test_transform_mat(self, device, dtype):
         aug = RandAugment(n=3, m=15)
         in_tensor = torch.rand(10, 3, 50, 50, device=device, dtype=dtype, requires_grad=True)