ivy pad backend consistency

ivy/array/extensions/layers.py

@@ -1,6 +1,6 @@
 # global
 import abc
-from typing import Optional, Union, Tuple, Iterable, Callable, Literal
+from typing import Optional, Union, Tuple, Iterable, Callable, Literal, Any
 from numbers import Number
 
 # local
@@ -157,8 +157,8 @@ def max_pool2d(
 
     def pad(
         self: ivy.Array,
-        /,
         pad_width: Union[Iterable[Tuple[int]], int],
+        /,
         *,
         mode: Optional[
             Union[
@@ -183,6 +183,7 @@ def pad(
         end_values: Optional[Union[Iterable[Tuple[Number]], Number]] = 0,
         reflect_type: Optional[Literal["even", "odd"]] = "even",
         out: Optional[ivy.Array] = None,
+        **kwargs: Optional[Any],
     ) -> ivy.Array:
         """
         ivy.Array instance method variant of ivy.pad. This method simply
@@ -198,4 +199,5 @@ def pad(
             end_values=end_values,
             reflect_type=reflect_type,
             out=out,
+            **kwargs,
         )

ivy/container/extensions/layers.py

@@ -8,6 +8,7 @@
     Callable,
     Literal,
     Iterable,
+    Any,
 )
 from numbers import Number
 
@@ -460,10 +461,10 @@ def kaiser_window(
 
     @staticmethod
     def static_pad(
-        x: ivy.Container,
+        input: ivy.Container,
+        pad_width: Union[Iterable[Tuple[int]], int],
         /,
         *,
-        pad_width: Union[Iterable[Tuple[int]], int],
         mode: Optional[
             Union[
                 Literal[
@@ -491,6 +492,7 @@ def static_pad(
         prune_unapplied: bool = False,
         map_sequences: bool = False,
         out: Optional[ivy.Container] = None,
+        **kwargs: Optional[Any],
     ) -> ivy.Container:
         """
         ivy.Container static method variant of ivy.pad. This method simply
@@ -499,7 +501,7 @@ def static_pad(
         """
         return ContainerBase.multi_map_in_static_method(
             "pad",
-            x,
+            input,
             pad_width,
             mode=mode,
             stat_length=stat_length,
@@ -511,13 +513,14 @@ def static_pad(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            **kwargs,
         )
 
     def pad(
         self: ivy.Container,
+        pad_width: Union[Iterable[Tuple[int]], int],
         /,
         *,
-        pad_width: Union[Iterable[Tuple[int]], int],
         mode: Optional[
             Union[
                 Literal[
@@ -545,6 +548,7 @@ def pad(
         prune_unapplied: bool = False,
         map_sequences: bool = False,
         out: Optional[ivy.Container] = None,
+        **kwargs: Optional[Any],
     ) -> ivy.Container:
         """
         ivy.Container instance method variant of ivy.pad. This method simply
@@ -564,4 +568,5 @@ def pad(
             prune_unapplied=prune_unapplied,
             map_sequences=map_sequences,
             out=out,
+            **kwargs,
         )

ivy/functional/backends/jax/extensions/layers.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Tuple, Callable, Literal, Sequence
+from typing import Optional, Union, Tuple, Callable, Literal, Sequence, Any
 from numbers import Number
 import ivy
 from ivy.functional.backends.jax import JaxArray
@@ -98,14 +98,28 @@ def max_pool2d(
     return res
 
 
+def kaiser_window(
+    window_length: int,
+    periodic: bool = True,
+    beta: float = 12.0,
+    *,
+    dtype: Optional[jnp.dtype] = None,
+    out: Optional[JaxArray] = None,
+) -> JaxArray:
+    if periodic is False:
+        return jnp.array(jnp.kaiser(M=window_length, beta=beta), dtype=dtype)
+    else:
+        return jnp.array(jnp.kaiser(M=window_length + 1, beta=beta)[:-1], dtype=dtype)
+
+
 def _flat_array_to_1_dim_array(x):
     return x.reshape((1,)) if x.shape == () else x
 
 
 def pad(
-    x: JaxArray,
-    /,
+    input: JaxArray,
     pad_width: Union[Sequence[Sequence[int]], JaxArray, int],
+    /,
     *,
     mode: Optional[
         Union[
@@ -130,52 +144,46 @@ def pad(
     end_values: Optional[Union[Sequence[Sequence[Number]], Number]] = 0,
     reflect_type: Optional[Literal["even", "odd"]] = "even",
     out: Optional[JaxArray] = None,
+    **kwargs: Optional[Any],
 ) -> JaxArray:
+    if callable(mode):
+        return jnp.pad(
+            _flat_array_to_1_dim_array(input),
+            pad_width,
+            mode=mode,
+            **kwargs,
+        )
     if mode in ["maximum", "mean", "median", "minimum"]:
         return jnp.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             stat_length=stat_length,
         )
     elif mode == "constant":
         return jnp.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             constant_values=constant_values,
         )
     elif mode == "linear_ramp":
         return jnp.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             end_values=end_values,
         )
     elif mode in ["reflect", "symmetric"]:
         return jnp.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             reflect_type=reflect_type,
         )
     else:
         return jnp.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
         )
-
-
-def kaiser_window(
-    window_length: int,
-    periodic: bool = True,
-    beta: float = 12.0,
-    *,
-    dtype: Optional[jnp.dtype] = None,
-    out: Optional[JaxArray] = None,
-) -> JaxArray:
-    if periodic is False:
-        return jnp.array(jnp.kaiser(M=window_length, beta=beta), dtype=dtype)
-    else:
-        return jnp.array(jnp.kaiser(M=window_length + 1, beta=beta)[:-1], dtype=dtype)

ivy/functional/backends/numpy/extensions/layers.py

@@ -2,7 +2,7 @@
 import math
 from numbers import Number
 import numpy as np
-from typing import Optional, Union, Tuple, Sequence, Callable, Literal
+from typing import Optional, Union, Tuple, Sequence, Callable, Literal, Any
 
 # local
 import ivy
@@ -107,14 +107,31 @@ def max_pool2d(
     return res
 
 
+def kaiser_window(
+    window_length: int,
+    periodic: bool = True,
+    beta: float = 12.0,
+    *,
+    dtype: Optional[np.dtype] = None,
+    out: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    if periodic is False:
+        return np.array(np.kaiser(M=window_length, beta=beta), dtype=dtype)
+    else:
+        return np.array(np.kaiser(M=window_length + 1, beta=beta)[:-1], dtype=dtype)
+
+
+kaiser_window.support_native_out = False
+
+
 def _flat_array_to_1_dim_array(x):
     return x.reshape((1,)) if x.shape == () else x
 
 
 def pad(
-    x: np.ndarray,
-    /,
+    input: np.ndarray,
     pad_width: Union[Sequence[Sequence[int]], np.ndarray, int],
+    /,
     *,
     mode: Optional[
         Union[
@@ -139,55 +156,46 @@ def pad(
     end_values: Optional[Union[Sequence[Sequence[Number]], Number]] = 0,
     reflect_type: Optional[Literal["even", "odd"]] = "even",
     out: Optional[np.ndarray] = None,
+    **kwargs: Optional[Any],
 ) -> np.ndarray:
+    if callable(mode):
+        return np.pad(
+            _flat_array_to_1_dim_array(input),
+            pad_width,
+            mode=mode,
+            **kwargs,
+        )
     if mode in ["maximum", "mean", "median", "minimum"]:
         return np.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             stat_length=stat_length,
         )
     elif mode == "constant":
         return np.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             constant_values=constant_values,
         )
     elif mode == "linear_ramp":
         return np.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             end_values=end_values,
         )
     elif mode in ["reflect", "symmetric"]:
         return np.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
             reflect_type=reflect_type,
         )
     else:
         return np.pad(
-            _flat_array_to_1_dim_array(x),
+            _flat_array_to_1_dim_array(input),
             pad_width,
             mode=mode,
         )
-
-
-def kaiser_window(
-    window_length: int,
-    periodic: bool = True,
-    beta: float = 12.0,
-    *,
-    dtype: Optional[np.dtype] = None,
-    out: Optional[np.ndarray] = None,
-) -> np.ndarray:
-    if periodic is False:
-        return np.array(np.kaiser(M=window_length, beta=beta), dtype=dtype)
-    else:
-        return np.array(np.kaiser(M=window_length + 1, beta=beta)[:-1], dtype=dtype)
-
-
-kaiser_window.support_native_out = False

ivy/functional/backends/tensorflow/extensions/layers.py

@@ -1,5 +1,4 @@
-from typing import Union, Optional, Tuple, Literal
-from numbers import Number
+from typing import Union, Optional, Tuple
 import tensorflow as tf
 
 
@@ -42,31 +41,6 @@ def max_pool2d(
     return res
 
 
-def pad(
-    x: tf.Tensor,
-    /,
-    pad_width: tf.Tensor,
-    *,
-    mode: Optional[Literal["constant", "reflect", "symmetric"]] = "constant",
-    stat_length: Optional[Union[tf.Tensor, int]] = None,
-    constant_values: Optional[Number] = 0,
-    end_values: Optional[Number] = 0,
-    reflect_type: Optional[Literal["even", "odd"]] = "even",
-    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
-) -> tf.Tensor:
-    if x.shape == ():
-        x = tf.reshape(x, (-1,))
-    if mode == "constant":
-        return tf.pad(
-            x,
-            pad_width,
-            mode=mode,
-            constant_values=constant_values,
-        )
-    else:
-        return tf.pad(x, pad_width, mode=mode)
-
-
 def kaiser_window(
     window_length: int,
     periodic: bool = True,