From c6ccf22a80d83acf9738f4e4c0d90cae6ed9f952 Mon Sep 17 00:00:00 2001 From: arshPratap Date: Fri, 25 Aug 2023 20:30:00 +0530 Subject: [PATCH] feat: Added adaptive max pooling 1D to Ivy backend --- ivy/data_classes/array/experimental/layers.py | 25 ++++++ .../container/experimental/layers.py | 68 ++++++++++++++++ ivy/functional/ivy/experimental/layers.py | 78 +++++++++++++++++++ .../test_experimental/test_nn/test_layers.py | 29 +++++++ 4 files changed, 200 insertions(+) diff --git a/ivy/data_classes/array/experimental/layers.py b/ivy/data_classes/array/experimental/layers.py index eefb667184fb4..88698f68d465b 100644 --- a/ivy/data_classes/array/experimental/layers.py +++ b/ivy/data_classes/array/experimental/layers.py @@ -1328,3 +1328,28 @@ def max_unpool1d( padding=padding, data_format=data_format, ) + + def adaptive_max_pool1d( + self: ivy.Array, + output_size: int, + ) -> ivy.Array: + """ + Apply a 1D adaptive maximum pooling over an input signal composed of several + input planes. + Parameters + ---------- + self + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. + Returns + ------- + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` + """ + return ivy.adaptive_max_pool1d( + self._data, + output_size, + ) diff --git a/ivy/data_classes/container/experimental/layers.py b/ivy/data_classes/container/experimental/layers.py index fc61b0caa19dd..a6d59ba416aeb 100644 --- a/ivy/data_classes/container/experimental/layers.py +++ b/ivy/data_classes/container/experimental/layers.py @@ -2854,3 +2854,71 @@ def max_unpool1d( padding=padding, data_format=data_format, ) + + @staticmethod + def static_adaptive_max_pool1d( + input: Union[ivy.Array, ivy.NativeArray, ivy.Container], + output_size: Union[int, ivy.Container], + *, + key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, + to_apply: Union[bool, ivy.Container] = True, + prune_unapplied: Union[bool, ivy.Container] = False, + map_sequences: Union[bool, ivy.Container] = False, + ) -> ivy.Container: + """ + ivy.Container static method variant of ivy.adaptive_max_pool1d. This method + simply wraps the function, and so the docstring for ivy.adaptive_max_pool1d also + applies to this method with minimal changes. + Parameters + ---------- + input + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. + Returns + ------- + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` + """ + return ContainerBase.cont_multi_map_in_function( + "adaptive_max_pool1d", + input, + output_size, + key_chains=key_chains, + to_apply=to_apply, + prune_unapplied=prune_unapplied, + map_sequences=map_sequences, + ) + + def adaptive_max_pool1d( + self: ivy.Container, + output_size: Union[int, ivy.Container], + *, + key_chains: Optional[Union[List[str], Dict[str, str], ivy.Container]] = None, + to_apply: Union[bool, ivy.Container] = True, + prune_unapplied: Union[bool, ivy.Container] = False, + map_sequences: Union[bool, ivy.Container] = False, + ) -> ivy.Container: + """ + Apply a 1D adaptive maximum pooling over an input signal composed of several + input planes. + Parameters + ---------- + self + Input container. + output_size + Spatial output size. + Returns + ------- + The result of the pooling operation. + """ + return self.static_adaptive_max_pool1d( + self, + output_size, + key_chains=key_chains, + to_apply=to_apply, + prune_unapplied=prune_unapplied, + map_sequences=map_sequences, + ) diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py index 2b8a3e9102014..6546104edcccd 100644 --- a/ivy/functional/ivy/experimental/layers.py +++ b/ivy/functional/ivy/experimental/layers.py @@ -3221,3 +3221,81 @@ def max_unpool1d( ), "to_skip": ("inputs_to_ivy_arrays", "handle_partial_mixed_function"), } + + +@handle_exceptions +@inputs_to_ivy_arrays +def adaptive_max_pool1d( + input: Union[ivy.Array, ivy.NativeArray], + output_size: int, +) -> ivy.Array: + """ + Apply a 1D adaptive maximum pooling over an input signal composed of several input + planes. + Parameters + ---------- + input + Input array. Must have shape (N, C, L_in) or (C, L_in) where N is + the batch dimension, C is the feature dimension, and L_in is the spatial + dimension. + output_size + Spatial output size. + Returns + ------- + The result of the pooling operation. Will have shape (N, C, L_out) or + (C, L_out), where L_out = `output_size` + """ + squeeze = False + if input.ndim == 2: + input = ivy.expand_dims(input, axis=0) + squeeze = True + elif input.ndim != 3: + raise ivy.utils.exceptions.IvyException( + f"Got {len(input.shape)}D input, but only 2D and 3D inputs are supported.", + ) + + if input.shape[-1] % output_size == 0: + stride = input.shape[-1] // output_size + kernel_size = input.shape[-1] - (output_size - 1) * stride + pooled_output = ivy.max_pool1d( + input, kernel_size, stride, "VALID", data_format="NCW" + ) + if squeeze: + return ivy.squeeze(pooled_output, axis=0) + return pooled_output + + idxw, length_w, range_max_w, adaptive_w = _compute_idx( + input.shape[-1], output_size, input.device + ) + + # to numpy and back in order to bypass a slicing error in tensorflow + vals = ivy.array(input.to_numpy()[..., idxw]) + + if not adaptive_w: + ret = ivy.max(vals, axis=-1) + ret = ivy.squeeze(ret, axis=0) if squeeze else ret + return ret + + vals, length_w = _mask(vals, length_w, range_max_w, dim=-1) + + ret = None + for i in range(vals.shape[-1]): + if ret is None: + ret = vals[..., i] + else: + ret = ivy.maximum(ret, vals[..., i]) + pooled_output = ret.astype(vals.dtype) + + pooled_output = ivy.squeeze(pooled_output, axis=0) if squeeze else pooled_output + return pooled_output + + +adaptive_max_pool1d.mixed_backend_wrappers = { + "to_add": ( + "handle_backend_invalid", + "inputs_to_native_arrays", + "outputs_to_ivy_arrays", + "handle_device_shifting", + ), + "to_skip": ("inputs_to_ivy_arrays",), +} diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py index 81eec1e5f3bf0..f1f35be208843 100644 --- a/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py +++ b/ivy_tests/test_ivy/test_functional/test_experimental/test_nn/test_layers.py @@ -546,6 +546,35 @@ def test_adaptive_avg_pool2d( ) +@handle_test( + fn_tree="functional.ivy.experimental.adaptive_max_pool1d", + dtype_and_x=helpers.dtype_and_values( + available_dtypes=helpers.get_dtypes("float"), + min_num_dims=2, + max_num_dims=3, + min_dim_size=1, + max_value=100, + min_value=-100, + ), + output_size=helpers.ints(min_value=1, max_value=5), + test_with_out=st.just(False), + ground_truth_backend="torch", +) +def test_adaptive_max_pool1d( + *, dtype_and_x, output_size, test_flags, backend_fw, fn_name, on_device +): + input_dtype, x = dtype_and_x + helpers.test_function( + input_dtypes=input_dtype, + test_flags=test_flags, + backend_to_test=backend_fw, + fn_name=fn_name, + on_device=on_device, + input=x[0], + output_size=output_size, + ) + + @handle_test( fn_tree="functional.ivy.experimental.adaptive_max_pool2d", dtype_and_x=helpers.dtype_and_values(