feat: reflection_padding layer added with test cases

Author: thesahibnanda

.vscode/settings.json

@@ -18,8 +18,8 @@
   ],
   "[python]": {
       "editor.codeActionsOnSave": {
-          "source.organizeImports": true
-      }
+        "source.organizeImports": "explicit"
+    }
   },
   "python.analysis.diagnosticSeverityOverrides": {
       "reportMissingImports": "none"

tf_keras/layers/convolutional/BUILD

@@ -215,86 +215,97 @@ py_library(
 )
 
 py_library(
-    name = "depthwise_conv1d",
-    srcs = ["depthwise_conv1d.py"],
-    srcs_version = "PY3",
-    deps = [
-        ":base_depthwise_conv",
-        "//:expect_tensorflow_installed",
-        "//tf_keras/utils:engine_utils",
-        "//tf_keras/utils:tf_utils",
-    ],
+name = "depthwise_conv1d",
+srcs = ["depthwise_conv1d.py"],
+srcs_version = "PY3",
+deps = [
+    ":base_depthwise_conv",
+    "//:expect_tensorflow_installed",
+    "//tf_keras/utils:engine_utils",
+    "//tf_keras/utils:tf_utils",
+],
 )
 
 py_library(
-    name = "depthwise_conv2d",
-    srcs = ["depthwise_conv2d.py"],
-    srcs_version = "PY3",
-    deps = [
-        ":base_depthwise_conv",
-        "//tf_keras:backend",
-        "//tf_keras/utils:engine_utils",
-        "//tf_keras/utils:tf_utils",
-    ],
+name = "depthwise_conv2d",
+srcs = ["depthwise_conv2d.py"],
+srcs_version = "PY3",
+deps = [
+    ":base_depthwise_conv",
+    "//tf_keras:backend",
+    "//tf_keras/utils:engine_utils",
+    "//tf_keras/utils:tf_utils",
+],
+)
+
+py_library(
+name = "reflection_padding",
+srcs = ["reflection_padding.py"],  # Adjust this to your actual file name
+srcs_version = "PY3",
+deps = [
+    "//:expect_tensorflow_installed",  # Assuming reflection_padding.py depends on TensorFlow
+    "//tf_keras/utils:engine_utils",
+    "//tf_keras/utils:tf_utils",
+],
 )
 
 cuda_py_test(
-    name = "conv_test",
-    size = "medium",
-    srcs = ["conv_test.py"],
-    python_version = "PY3",
-    shard_count = 8,
-    deps = [
-        "//:expect_absl_installed",  # absl/testing:parameterized
-        "//:expect_numpy_installed",
-        "//:expect_tensorflow_installed",
-        "//tf_keras",
-        "//tf_keras/testing_infra:test_combinations",
-        "//tf_keras/testing_infra:test_utils",
-    ],
+name = "conv_test",
+size = "medium",
+srcs = ["conv_test.py"],
+python_version = "PY3",
+shard_count = 8,
+deps = [
+    "//:expect_absl_installed",  # absl/testing:parameterized
+    "//:expect_numpy_installed",
+    "//:expect_tensorflow_installed",
+    "//tf_keras",
+    "//tf_keras/testing_infra:test_combinations",
+    "//tf_keras/testing_infra:test_utils",
+],
 )
 
 cuda_py_test(
-    name = "conv_transpose_test",
-    size = "medium",
-    srcs = ["conv_transpose_test.py"],
-    python_version = "PY3",
-    deps = [
-        "//:expect_absl_installed",  # absl/testing:parameterized
-        "//:expect_numpy_installed",
-        "//:expect_tensorflow_installed",
-        "//tf_keras",
-        "//tf_keras/testing_infra:test_combinations",
-        "//tf_keras/testing_infra:test_utils",
-    ],
+name = "conv_transpose_test",
+size = "medium",
+srcs = ["conv_transpose_test.py"],
+python_version = "PY3",
+deps = [
+    "//:expect_absl_installed",  # absl/testing:parameterized
+    "//:expect_numpy_installed",
+    "//:expect_tensorflow_installed",
+    "//tf_keras",
+    "//tf_keras/testing_infra:test_combinations",
+    "//tf_keras/testing_infra:test_utils",
+],
 )
 
 cuda_py_test(
-    name = "depthwise_conv_test",
-    size = "medium",
-    srcs = ["depthwise_conv_test.py"],
-    python_version = "PY3",
-    shard_count = 8,
-    deps = [
-        "//:expect_absl_installed",  # absl/testing:parameterized
-        "//:expect_tensorflow_installed",
-        "//tf_keras",
-        "//tf_keras/testing_infra:test_combinations",
-        "//tf_keras/testing_infra:test_utils",
-    ],
+name = "depthwise_conv_test",
+size = "medium",
+srcs = ["depthwise_conv_test.py"],
+python_version = "PY3",
+shard_count = 8,
+deps = [
+    "//:expect_absl_installed",  # absl/testing:parameterized
+    "//:expect_tensorflow_installed",
+    "//tf_keras",
+    "//tf_keras/testing_infra:test_combinations",
+    "//tf_keras/testing_infra:test_utils",
+],
 )
 
 cuda_py_test(
-    name = "separable_conv_test",
-    size = "medium",
-    srcs = ["separable_conv_test.py"],
-    python_version = "PY3",
-    deps = [
-        "//:expect_absl_installed",  # absl/testing:parameterized
-        "//:expect_numpy_installed",
-        "//:expect_tensorflow_installed",
-        "//tf_keras",
-        "//tf_keras/testing_infra:test_combinations",
-        "//tf_keras/testing_infra:test_utils",
-    ],
+name = "separable_conv_test",
+size = "medium",
+srcs = ["separable_conv_test.py"],
+python_version = "PY3",
+deps = [
+    "//:expect_absl_installed",  # absl/testing:parameterized
+    "//:expect_numpy_installed",
+    "//:expect_tensorflow_installed",
+    "//tf_keras",
+    "//tf_keras/testing_infra:test_combinations",
+    "//tf_keras/testing_infra:test_utils",
+],
 )

tf_keras/layers/convolutional/base_reflection_padding.py

@@ -0,0 +1,53 @@
+import tensorflow as tf
+from tensorflow.keras.layers import Layer
+
+
+class BaseReflectionPadding(Layer):
+    """Abstract N-D reflection padding layer.
+
+    This layer performs reflection padding on the input tensor.
+
+    Args:
+        padding: int, or tuple/list of n ints, for n > 1.
+            If int: the same symmetric padding is applied to all spatial dimensions.
+            If tuple/list of n ints: interpreted as n different symmetric padding values
+            for each spatial dimension. 
+            No padding is applied to the batch or channel dimensions.
+
+    Raises:
+        ValueError: If `padding` is negative or not of length 2 or more.
+
+    """
+
+    def __init__(self, padding=(1, 1), **kwargs):
+        super(BaseReflectionPadding, self).__init__(**kwargs)
+        if isinstance(padding, int):
+            self.padding = (padding, padding)
+        elif isinstance(padding, tuple) or isinstance(padding, list):
+            if len(padding) != self.rank:
+                raise ValueError(
+                    f"If passing a tuple or list as padding, it must be of length {self.rank}. Received length: {len(padding)}"
+                )
+            self.padding = padding
+        else:
+            raise ValueError(
+                f"Unsupported padding type. Expected int, tuple, or list. Received: {type(padding)}"
+            )
+
+        for pad in self.padding:
+            if pad < 0:
+                raise ValueError("Padding cannot be negative.")
+
+    def compute_output_shape(self, input_shape):
+        output_shape = list(input_shape)
+        for i in range(1, self.rank + 1):
+            output_shape[i] += 2 * self.padding[i - 1]
+        return tuple(output_shape)
+
+    def call(self, inputs):
+        padding_dims = [[0, 0]]
+        for pad in self.padding:
+            padding_dims.append([pad, pad])
+        for _ in range(self.rank - len(self.padding)):
+            padding_dims.append([0, 0])
+        return tf.pad(inputs, padding_dims, mode='REFLECT')

tf_keras/layers/convolutional/reflection_padding1d.py

@@ -0,0 +1,20 @@
+from base_reflection_padding import BaseReflectionPadding
+
+
+class ReflectionPadding1D(BaseReflectionPadding):
+    """1D reflection padding layer.
+
+    This layer performs reflection padding on a 1D input tensor.
+    Inherits from BaseReflectionPadding.
+
+    Args:
+        padding: int, or tuple/list of 1 int, specifying the padding width.
+            If int: the same symmetric padding is applied to both sides.
+            If tuple/list of 1 int: interpreted as two different symmetric padding values.
+    """
+
+    rank = 1
+
+    def __init__(self, padding=1, **kwargs):
+        super(ReflectionPadding1D, self).__init__(
+            padding=(padding,), **kwargs)

tf_keras/layers/convolutional/reflection_padding2d.py

@@ -0,0 +1,21 @@
+from base_reflection_padding import BaseReflectionPadding
+
+
+class ReflectionPadding2D(BaseReflectionPadding):
+    """2D reflection padding layer.
+
+    This layer performs reflection padding on a 2D input tensor.
+    Inherits from BaseReflectionPadding.
+
+    Args:
+        padding: int, or tuple/list of 2 ints, for height and width respectively.
+            If int: the same symmetric padding is applied to both dimensions.
+            If tuple/list of 2 ints: interpreted as two different symmetric padding values.
+    """
+
+    rank = 2
+
+    def __init__(self, padding=(1, 1), **kwargs):
+        if isinstance(padding, int):
+            padding = (padding, padding)
+        super(ReflectionPadding2D, self).__init__(padding=padding, **kwargs)

tf_keras/layers/convolutional/reflection_padding3d.py

@@ -0,0 +1,21 @@
+from base_reflection_padding import BaseReflectionPadding
+
+
+class ReflectionPadding3D(BaseReflectionPadding):
+    """3D reflection padding layer.
+
+    This layer performs reflection padding on a 3D input tensor.
+    Inherits from BaseReflectionPadding.
+
+    Args:
+        padding: int, or tuple/list of 3 ints, for height, width, and depth respectively.
+            If int: the same symmetric padding is applied to all dimensions.
+            If tuple/list of 3 ints: interpreted as three different symmetric padding values.
+    """
+
+    rank = 3
+
+    def __init__(self, padding=(1, 1, 1), **kwargs):
+        if isinstance(padding, int) and padding > 0:
+            padding = (padding, padding, padding)
+        super(ReflectionPadding3D, self).__init__(padding=padding, **kwargs)

tf_keras/layers/convolutional/reflection_padding_test.py

@@ -0,0 +1,46 @@
+import numpy as np
+import tensorflow as tf
+from absl.testing import parameterized
+from reflection_padding1d import ReflectionPadding1D
+from reflection_padding2d import ReflectionPadding2D
+from reflection_padding3d import ReflectionPadding3D
+
+from tf_keras.testing_infra import test_combinations
+from tf_keras.testing_infra import test_utils
+
+
+@test_combinations.run_all_keras_modes
+class ReflectionPaddingTest(test_combinations.TestCase):
+    def _run_test(self, padding_layer_cls, kwargs, input_shape, expected_output_shape):
+        with self.cached_session():
+            test_utils.layer_test(
+                padding_layer_cls,
+                kwargs=kwargs,
+                input_shape=input_shape,
+                expected_output_shape=expected_output_shape,
+            )
+
+    @parameterized.named_parameters(
+        ("ReflectionPadding1D", ReflectionPadding1D,
+         {"padding": 2}, (None, 5, 3), (None, 9, 3)),
+        ("ReflectionPadding2D", ReflectionPadding2D, {
+         "padding": (2, 1)}, (None, 5, 6, 3), (None, 9, 8, 3)),
+        ("ReflectionPadding3D", ReflectionPadding3D, {
+         "padding": (1, 2, 3)}, (None, 5, 6, 7, 3), (None, 7, 10, 13, 3)),
+    )
+    def test_reflection_padding(self, padding_layer_cls, kwargs, input_shape, expected_output_shape):
+        self._run_test(padding_layer_cls, kwargs,
+                       input_shape, expected_output_shape)
+
+    def test_reflection_padding_dynamic_shape(self):
+        with self.cached_session():
+            layer = ReflectionPadding2D(padding=(2, 2))
+            input_shape = (None, None, None, 3)
+            inputs = tf.keras.Input(shape=input_shape)
+            x = layer(inputs)
+            # Won't raise error here with None values in input shape.
+            layer(x)
+
+
+if __name__ == "__main__":
+    tf.test.main()