Repository: ivy-llc/ivy

Issue:
channel_shuffle



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/paddle/nn/functional/vision.py b/ivy/functional/frontends/paddle/nn/functional/vision.py
--- a/ivy/functional/frontends/paddle/nn/functional/vision.py
+++ b/ivy/functional/frontends/paddle/nn/functional/vision.py
@@ -1,7 +1,7 @@
 # local
 
 import ivy
-from ivy.func_wrapper import with_unsupported_dtypes
+from ivy.func_wrapper import with_unsupported_dtypes, with_supported_dtypes
 from ivy.functional.frontends.paddle.func_wrapper import (
     to_ivy_arrays_and_back,
 )
@@ -135,3 +135,38 @@ def affine_grid(theta, out_shape, align_corners=True):
             base_grid[:, :, :, :, 3] = ivy.full((D, H, W), 1)
             grid = ivy.matmul(base_grid.view((N, D * H * W, 4)), theta.swapaxes(1, 2))
             return grid.view((N, D, H, W, 3))
+
+
+@to_ivy_arrays_and_back
+@with_supported_dtypes({"2.5.1 and below": ("float32", "float64")}, "paddle")
+def channel_shuffle(x, groups, data_format="NCHW", name=None):
+    if len(ivy.shape(x)) != 4:
+        raise ValueError(
+            "Input x should be 4D tensor, but received x with the shape of {}".format(
+                ivy.shape(x)
+            )
+        )
+
+    if not isinstance(groups, int):
+        raise TypeError("groups must be int type")
+
+    if groups <= 0:
+        raise ValueError("groups must be positive")
+
+    if data_format not in ["NCHW", "NHWC"]:
+        raise ValueError(
+            "Attr(data_format) should be 'NCHW' or 'NHWC'."
+            "But recevie Attr(data_format): {} ".format(data_format)
+        )
+
+    if data_format == "NCHW":
+        b, c, h, w = ivy.shape(x)
+        x = ivy.reshape(x, (b, groups, c // groups, h, w))
+        x = ivy.permute_dims(x, (0, 2, 1, 3, 4))
+        x = ivy.reshape(x, (b, c, h, w))
+    else:
+        b, h, w, c = ivy.shape(x)
+        x = ivy.reshape(x, (b, h, w, groups, c // groups))
+        x = ivy.permute_dims(x, (0, 1, 2, 4, 3))
+        x = ivy.reshape(x, (b, h, w, c))
+    return x


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_paddle/test_nn/test_functional/test_vision.py b/ivy_tests/test_ivy/test_frontends/test_paddle/test_nn/test_functional/test_vision.py
--- a/ivy_tests/test_ivy/test_frontends/test_paddle/test_nn/test_functional/test_vision.py
+++ b/ivy_tests/test_ivy/test_frontends/test_paddle/test_nn/test_functional/test_vision.py
@@ -112,3 +112,57 @@ def test_paddle_affine_grid(
         out_shape=size,
         align_corners=align_corners,
     )
+
+
+@st.composite
+def _image_shape_helper(draw, data_format):
+    n = draw(helpers.ints(min_value=1, max_value=10), label="batch")
+    c = draw(st.sampled_from([1, 3]), label="channel")
+    h = draw(helpers.ints(min_value=1, max_value=100), label="height")
+    w = draw(helpers.ints(min_value=1, max_value=100), label="width")
+
+    if data_format == "NCHW":
+        shape = (n, c, h, w)
+    else:
+        shape = (n, h, w, c)
+
+    return shape
+
+
+# channel_shuffle
+@handle_frontend_test(
+    fn_tree="paddle.nn.functional.channel_shuffle",
+    dtype_and_x=helpers.dtype_and_values(
+        available_dtypes=["float32", "float64"],
+        shape=_image_shape_helper(data_format=st.sampled_from(["NCHW", "NHWC"])),
+    ),
+    groups=helpers.ints(min_value=1),
+    data_format=st.sampled_from(["NCHW", "NHWC"]),
+)
+def test_paddle_channel_shuffle(
+    *,
+    dtype_and_x,
+    groups,
+    data_format,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+    backend_fw,
+):
+    input_dtype, x = dtype_and_x
+    if data_format == "NCHW":
+        assume(ivy.shape(x[0])[1] % groups == 0)
+    else:
+        assume(ivy.shape(x[0])[3] % groups == 0)
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        test_flags=test_flags,
+        backend_to_test=backend_fw,
+        on_device=on_device,
+        frontend=frontend,
+        fn_tree=fn_tree,
+        x=x[0],
+        groups=groups,
+        data_format=data_format,
+    )