Repository: kornia/kornia

Issue:
Add/move `_differentiable_clipping`, ` _differentiable_polynomial_floor` and ` _differentiable_polynomial_rounding` to our utilities functions
              for a further pr, i think having `_differentiable_clipping`, ` _differentiable_polynomial_floor` and ` _differentiable_polynomial_rounding` into our utilities can be usefully :)

_Originally posted by @johnnv1 in https://github.com/kornia/kornia/pull/2748#discussion_r1464089248_
            


Discussion:
@johnnv1 Where do you prefer moving these in utils ?  `utils/helpers.py` a good option ?

Resolving patch:
diff --git a/kornia/enhance/jpeg.py b/kornia/enhance/jpeg.py
--- a/kornia/enhance/jpeg.py
+++ b/kornia/enhance/jpeg.py
@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import Optional
-
 import torch
 
 from kornia.color import rgb_to_ycbcr, ycbcr_to_rgb
@@ -14,6 +12,11 @@
 )
 from kornia.geometry.transform.affwarp import rescale
 from kornia.utils.image import perform_keep_shape_image
+from kornia.utils.misc import (
+    differentiable_clipping,
+    differentiable_polynomial_floor,
+    differentiable_polynomial_rounding,
+)
 
 __all__ = ["jpeg_codec_differentiable", "JPEGCodecDifferentiable"]
 
@@ -136,61 +139,6 @@ def _idct_8x8(input: Tensor) -> Tensor:
     return output
 
 
-def _differentiable_polynomial_rounding(input: Tensor) -> Tensor:
-    """This function implements differentiable rounding.
-
-    Args:
-        input (Tensor): Input tensor of any shape to be rounded.
-
-    Returns:
-        output (Tensor): Pseudo rounded tensor of the same shape as input tensor.
-    """
-    # Perform differentiable rounding
-    output: Tensor = torch.round(input) + (input - torch.round(input)) ** 3
-    return output
-
-
-def _differentiable_polynomial_floor(input: Tensor) -> Tensor:
-    """This function implements differentiable floor.
-
-    Args:
-        input (Tensor): Input tensor of any shape to be floored.
-
-    Returns:
-        output (Tensor): Pseudo rounded tensor of the same shape as input tensor.
-    """
-    # Perform differentiable rounding
-    output: Tensor = torch.floor(input) + (input - 0.5 - torch.floor(input)) ** 3
-    return output
-
-
-def _differentiable_clipping(
-    input: Tensor,
-    min: Optional[float] = None,
-    max: Optional[float] = None,
-    scale: float = 0.02,
-) -> Tensor:
-    """This function implements a differentiable and soft approximation of the clipping operation.
-
-    Args:
-        input (Tensor): Input tensor of any shape.
-        min (Optional[float]): Minimum value.
-        max (Optional[float]): Maximum value.
-        scale (float): Scale value. Default 0.02.
-
-    Returns:
-        output (Tensor): Clipped output tensor of the same shape as the input tensor.
-    """
-    # Make a copy of the input tensor
-    output: Tensor = input.clone()
-    # Perform differentiable soft clipping
-    if max is not None:
-        output[output > max] = -scale * (torch.exp(-output[output > max] + max) - 1.0) + max
-    if min is not None:
-        output[output < min] = scale * (torch.exp(output[output < min] - min) - 1.0) + min
-    return output
-
-
 def _jpeg_quality_to_scale(
     compression_strength: Tensor,
 ) -> Tensor:
@@ -203,7 +151,7 @@ def _jpeg_quality_to_scale(
         scale (Tensor): Scaling factor to be applied to quantization matrix. Same shape as input.
     """
     # Get scale
-    scale: Tensor = _differentiable_polynomial_floor(
+    scale: Tensor = differentiable_polynomial_floor(
         torch.where(compression_strength < 50, 5000.0 / compression_strength, 200.0 - 2.0 * compression_strength)
     )
     return scale
@@ -229,12 +177,12 @@ def _quantize(
         quantization_table[:, None] * _jpeg_quality_to_scale(jpeg_quality)[:, None, None, None]
     )
     # Perform scaling
-    quantization_table = _differentiable_polynomial_floor(
-        _differentiable_clipping((quantization_table_scaled + 50.0) / 100.0, 1, 255)
+    quantization_table = differentiable_polynomial_floor(
+        differentiable_clipping((quantization_table_scaled + 50.0) / 100.0, 1, 255)
     )
     output: Tensor = input / quantization_table
     # Perform rounding
-    output = _differentiable_polynomial_rounding(output)
+    output = differentiable_polynomial_rounding(output)
     return output
 
 
@@ -258,8 +206,8 @@ def _dequantize(
         quantization_table[:, None] * _jpeg_quality_to_scale(jpeg_quality)[:, None, None, None]
     )
     # Perform scaling
-    output: Tensor = input * _differentiable_polynomial_floor(
-        _differentiable_clipping((quantization_table_scaled + 50.0) / 100.0, 1, 255)
+    output: Tensor = input * differentiable_polynomial_floor(
+        differentiable_clipping((quantization_table_scaled + 50.0) / 100.0, 1, 255)
     )
     return output
 
@@ -550,7 +498,7 @@ def jpeg_codec_differentiable(
         quantization_table_y=quantization_table_y,
     )
     # Clip coded image
-    image_rgb_jpeg = _differentiable_clipping(input=image_rgb_jpeg, min=0.0, max=255.0)
+    image_rgb_jpeg = differentiable_clipping(input=image_rgb_jpeg, min_val=0.0, max_val=255.0)
     # Back to original shape
     # image_rgb_jpeg = image_rgb_jpeg.view(original_shape)
     return image_rgb_jpeg
diff --git a/kornia/utils/__init__.py b/kornia/utils/__init__.py
--- a/kornia/utils/__init__.py
+++ b/kornia/utils/__init__.py
@@ -17,7 +17,13 @@
 from .image import ImageToTensor, image_list_to_tensor, image_to_tensor, tensor_to_image
 from .image_print import image_to_string, print_image
 from .memory import batched_forward
-from .misc import eye_like, vec_like
+from .misc import (
+    differentiable_clipping,
+    differentiable_polynomial_floor,
+    differentiable_polynomial_rounding,
+    eye_like,
+    vec_like,
+)
 from .one_hot import one_hot
 from .pointcloud_io import load_pointcloud_ply, save_pointcloud_ply
 
diff --git a/kornia/utils/misc.py b/kornia/utils/misc.py
--- a/kornia/utils/misc.py
+++ b/kornia/utils/misc.py
@@ -1,3 +1,7 @@
+from typing import Optional
+
+import torch
+
 from kornia.core import Tensor, eye, zeros
 
 
@@ -52,3 +56,60 @@ def vec_like(n: int, tensor: Tensor, shared_memory: bool = False) -> Tensor:
 
     vec = zeros(n, 1, device=tensor.device, dtype=tensor.dtype)
     return vec[None].expand(tensor.shape[0], n, 1) if shared_memory else vec[None].repeat(tensor.shape[0], 1, 1)
+
+
+def differentiable_polynomial_rounding(input: Tensor) -> Tensor:
+    """This function implements differentiable rounding.
+
+    Args:
+        input (Tensor): Input tensor of any shape to be rounded.
+
+    Returns:
+        output (Tensor): Pseudo rounded tensor of the same shape as input tensor.
+    """
+    # Perform differentiable rounding
+    input_round = input.round()
+    output: Tensor = input_round + (input - input_round) ** 3
+    return output
+
+
+def differentiable_polynomial_floor(input: Tensor) -> Tensor:
+    """This function implements differentiable floor.
+
+    Args:
+        input (Tensor): Input tensor of any shape to be floored.
+
+    Returns:
+        output (Tensor): Pseudo rounded tensor of the same shape as input tensor.
+    """
+    # Perform differentiable rounding
+    input_floor = input.floor()
+    output: Tensor = input_floor + (input - 0.5 - input_floor) ** 3
+    return output
+
+
+def differentiable_clipping(
+    input: Tensor,
+    min_val: Optional[float] = None,
+    max_val: Optional[float] = None,
+    scale: float = 0.02,
+) -> Tensor:
+    """This function implements a differentiable and soft approximation of the clipping operation.
+
+    Args:
+        input (Tensor): Input tensor of any shape.
+        min_val (Optional[float]): Minimum value.
+        max_val (Optional[float]): Maximum value.
+        scale (float): Scale value. Default 0.02.
+
+    Returns:
+        output (Tensor): Clipped output tensor of the same shape as the input tensor.
+    """
+    # Make a copy of the input tensor
+    output: Tensor = input.clone()
+    # Perform differentiable soft clipping
+    if max_val is not None:
+        output[output > max_val] = -scale * (torch.exp(-output[output > max_val] + max_val) - 1.0) + max_val
+    if min_val is not None:
+        output[output < min_val] = scale * (torch.exp(output[output < min_val] - min_val) - 1.0) + min_val
+    return output


Test patch:
diff --git a/tests/utils/test_misc.py b/tests/utils/test_misc.py
new file mode 100644
--- /dev/null
+++ b/tests/utils/test_misc.py
@@ -0,0 +1,47 @@
+from kornia.core import tensor
+from kornia.utils.misc import (
+    differentiable_clipping,
+    differentiable_polynomial_floor,
+    differentiable_polynomial_rounding,
+)
+
+from testing.base import BaseTester
+
+
+class TestDifferentiableClipping(BaseTester):
+    def test_differentiable_clipping(self, device):
+        x = tensor([1.0, 6.0, 10.0, 12.0], device=device)
+        y = differentiable_clipping(x, min_val=5.0, max_val=10.0)
+        y_expected = tensor([4.9804, 6.0, 10.0, 10.0173], device=device)
+
+        self.assert_close(y, y_expected)
+
+    def test_gradcheck(self, device):
+        x = tensor([1.0, 6.0, 11.0, 12.0], device=device)
+        self.gradcheck(differentiable_clipping, (x, 5.0, 10.0))
+
+
+class TestDifferentiablePolynomialRounding(BaseTester):
+    def test_differentiable_polynomial_rounding(self, device):
+        x = tensor([1.5], device=device)
+        y = differentiable_polynomial_rounding(x)
+        y_expected = tensor([1.875], device=device)
+
+        self.assert_close(y, y_expected)
+
+    def test_gradcheck(self, device):
+        x = tensor([1.0, 6.0, 10.0, 12.0], device=device)
+        self.gradcheck(differentiable_polynomial_rounding, (x))
+
+
+class TestDifferentiablePolynomialFloor(BaseTester):
+    def test_differentiable_polynomial_floor(self, device):
+        x = tensor([1.5], device=device)
+        y = differentiable_polynomial_floor(x)
+        y_expected = tensor([1.0], device=device)
+
+        self.assert_close(y, y_expected)
+
+    def test_gradcheck(self, device):
+        x = tensor([1.5, 3.1, 5.9, 6.6], device=device)
+        self.gradcheck(differentiable_polynomial_floor, (x))