import torch
import torch.nn as nn
import torch.nn.functional as F

N, C, H, W = 32, 64, 56, 56
GROUPS = 4


class ChannelShuffle(nn.Module):

    def __init__(self, groups):
        super().__init__()
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, num_channels, height, width = x.size()
        channels_per_group = num_channels // self.groups

        # 1. Reshape
        x = x.view(batch_size, self.groups, channels_per_group, height, width)

        # 2. Transpose (交换 groups 和 channels_per_group 维度)
        # dim 1 is groups, dim 2 is channels_per_group
        x = torch.transpose(x, 1, 2).contiguous()

        # 3. Flatten
        x = x.view(batch_size, num_channels, height, width)

        return x


class Model(nn.Module):
    def __init__(self, groups=GROUPS):
        super().__init__()
        self.op = ChannelShuffle(groups)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op(x)


def get_inputs():
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [GROUPS]