Repository: ivy-llc/ivy

Issue:
max_unpool2d
#3856


Discussion:

Resolving patch:
diff --git a/ivy/functional/backends/torch/experimental/layers.py b/ivy/functional/backends/torch/experimental/layers.py
--- a/ivy/functional/backends/torch/experimental/layers.py
+++ b/ivy/functional/backends/torch/experimental/layers.py
@@ -173,6 +173,69 @@ def max_pool2d(
     return res
 
 
+def max_unpool2d(
+    x: torch.Tensor,
+    indices: torch.Tensor,
+    kernel: Union[int, Tuple[int, ...]],
+    strides: Union[int, Tuple[int, ...]],
+    padding: Union[str, int, Tuple[int], List[Tuple[int, int]]],
+    /,
+    *,
+    data_format: str = "NHWC",
+    output_size: Optional[Tuple[int, ...]] = None,
+) -> torch.Tensor:
+    dims = 2
+
+    if data_format == "NHWC":
+        x = x.permute(0, 3, 1, 2)
+        indices = indices.permute(0, 3, 1, 2)
+        kernel = (
+            [kernel[i] for i in [0, 3, 1, 2]] if len(kernel) == (dims + 2) else kernel
+        )
+        strides = (
+            [strides[i] for i in [0, 3, 1, 2]]
+            if len(strides) == (dims + 2)
+            else strides
+        )
+        padding = (
+            [padding[i] for i in [0, 3, 1, 2]]
+            if isinstance(padding, list) and len(padding) == (dims + 2)
+            else padding
+        )
+
+    x, kernel, strides, depth_pooling = _determine_depth_max_pooling(
+        x, kernel, strides, dims, data_format="channel_first"
+    )
+    x_shape = list(x.shape[2:])
+    if not depth_pooling:
+        new_kernel = [kernel[i] + (kernel[i] - 1) * (1 - 1) for i in range(2)]
+
+        if isinstance(padding, str):
+            pad_h = _handle_padding(x_shape[0], strides[0], new_kernel[0], padding)
+            pad_w = _handle_padding(x_shape[1], strides[1], new_kernel[1], padding)
+            pad_list = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
+        else:
+            padding = (padding[1], padding[0])
+            pad_list = [item for sublist in padding for item in sublist]
+
+        unpooled_tensor = torch.nn.functional.max_unpool2d(
+            x,
+            indices,
+            kernel_size=kernel,
+            stride=strides,
+            padding=pad_list,
+            output_size=output_size,
+        )
+    else:
+        raise NotImplementedError(
+            "Unpooling for depthwise max pooling is not supported."
+        )
+
+    if data_format == "NHWC":
+        return unpooled_tensor.permute(0, 2, 3, 1)
+    return unpooled_tensor
+
+
 @with_unsupported_dtypes(
     {
         "2.0.1 and below": (
diff --git a/ivy/functional/frontends/torch/nn/functional/pooling_functions.py b/ivy/functional/frontends/torch/nn/functional/pooling_functions.py
--- a/ivy/functional/frontends/torch/nn/functional/pooling_functions.py
+++ b/ivy/functional/frontends/torch/nn/functional/pooling_functions.py
@@ -289,3 +289,12 @@ def avg_pool3d(
         count_include_pad=count_include_pad,
         divisor_override=divisor_override,
     )
+
+
+@to_ivy_arrays_and_back
+def max_unpool2d(input, indices, kernel_size, stride=None, padding=0, output_size=None):
+    if input.ndim == 3:
+        input = input.expand_dims()
+    if not stride:
+        stride = kernel_size
+    return ivy.max_unpool2d(input, indices, kernel_size, stride, padding, output_size)
diff --git a/ivy/functional/ivy/experimental/layers.py b/ivy/functional/ivy/experimental/layers.py
--- a/ivy/functional/ivy/experimental/layers.py
+++ b/ivy/functional/ivy/experimental/layers.py
@@ -354,6 +354,56 @@ def max_pool3d(
     )
 
 
+@handle_nestable
+@handle_out_argument
+@to_native_arrays_and_back
+@handle_device_shifting
+def max_unpool2d(
+    x: Union[ivy.Array, ivy.NativeArray],
+    kernel: Union[int, Tuple[int, ...]],
+    strides: Union[int, Tuple[int, ...]],
+    padding: Union[str, int, Tuple[int], List[Tuple[int, int]]],
+    /,
+    *,
+    out: Optional[ivy.Array] = None,
+) -> ivy.Array:
+    """
+    Compute a 2-D max unpool.
+
+    Parameters
+    ----------
+    x
+        Input image *[batch_size,h,w,d_in]*.
+    kernel
+        Size of the kernel i.e., the sliding window for each
+        dimension of input. *[h,w]*.
+    strides
+        The stride of the sliding window for each dimension of input.
+    padding
+        SAME" or "VALID" indicating the algorithm, or list
+        indicating the per-dimension paddings.
+    out
+        optional output array, for writing the result to.
+
+    Returns
+    -------
+    ret
+        The result of the pooling operation.
+
+    Both the description and the type hints above assumes an array input
+    for simplicity, but this function is *nestable*, and therefore
+    also accepts :class:`ivy.Container` instances in place of any of
+    the arguments.
+    """
+    return ivy.current_backend(x).max_unpool2d(
+        x,
+        kernel,
+        strides,
+        padding,
+        out=out,
+    )
+
+
 @handle_nestable
 @handle_out_argument
 @to_native_arrays_and_back


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_pooling_functions.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_pooling_functions.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_pooling_functions.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_nn/test_functional/test_pooling_functions.py
@@ -292,6 +292,61 @@ def test_torch_max_pool2d(
     )
 
 
+# max_unpool2d
+@handle_frontend_test(
+    fn_tree="torch.nn.functional.max_unpool2d",
+    x_k_indices_k_s_p=helpers.arrays_for_pooling(
+        min_dims=4,
+        max_dims=4,
+        min_side=1,
+        max_side=4,
+        explicit_or_str_padding=True,
+        return_indices=True,
+    ).filter(lambda x: x[4] != "VALID" and x[4] != "SAME"),
+    test_with_out=st.just(False),
+)
+def test_torch_max_unpool2d(
+    *,
+    x_k_indices_kernel_stride_padding,
+    test_flags,
+    frontend,
+    backend_fw,
+    fn_tree,
+    on_device,
+):
+    dtype, x, indices, kernel, stride, pad = x_k_indices_kernel_stride_padding
+
+    # Reshape the input data and indices as needed
+    x[0] = x[0].reshape((x[0].shape[0], x[0].shape[-1], *x[0].shape[1:-1]))
+    indices[0] = indices[0].reshape(
+        (indices[0].shape[0], indices[0].shape[-1], *indices[0].shape[1:-1])
+    )
+
+    # Calculate the output shape based on the input shape, kernel size, stride, and
+    # padding
+    output_shape = tuple(
+        [
+            (x[0].shape[i + 2] - 1) * stride[i] - 2 * pad[i] + kernel[i]
+            for i in range(len(kernel))
+        ]
+    )
+
+    helpers.test_frontend_function(
+        input_dtypes=dtype,
+        backend_to_test=backend_fw,
+        test_flags=test_flags,
+        frontend=frontend,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        input=x[0],
+        indices=indices[0],
+        output_size=output_shape,
+        kernel_size=kernel,
+        stride=stride,
+        padding=pad,
+    )
+
+
 # adaptive_avg_pool1d
 @handle_frontend_test(
     fn_tree="torch.nn.functional.adaptive_avg_pool1d",