L iJddlmcmZddlmZddlmZdgZGddeZ y)N)Tensor)ModuleChannelShufflecZeZdZUdZdgZeed<deddffd ZdedefdZ de fdZ xZ S) ra'Divides and rearranges the channels in a tensor. This operation divides the channels in a tensor of shape :math:`(N, C, *)` into g groups as :math:`(N, \frac{C}{g}, g, *)` and shuffles them, while retaining the original tensor shape in the final output. Args: groups (int): number of groups to divide channels in. Examples:: >>> channel_shuffle = nn.ChannelShuffle(2) >>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2) >>> input tensor([[[[ 1., 2.], [ 3., 4.]], [[ 5., 6.], [ 7., 8.]], [[ 9., 10.], [11., 12.]], [[13., 14.], [15., 16.]]]]) >>> output = channel_shuffle(input) >>> output tensor([[[[ 1., 2.], [ 3., 4.]], [[ 9., 10.], [11., 12.]], [[ 5., 6.], [ 7., 8.]], [[13., 14.], [15., 16.]]]]) groupsreturnNc0t|||_y)N)super__init__r)selfr __class__s e/mnt/ssd/data/python-lab/Trading/venv/lib/python3.12/site-packages/torch/nn/modules/channelshuffle.pyr zChannelShuffle.__init__0s  inputcBtj||jS)z( Runs the forward pass. )Fchannel_shuffler)r rs rforwardzChannelShuffle.forward4s   44rc d|jS)z@ Return the extra representation of the module. zgroups=)r)r s r extra_reprzChannelShuffle.extra_repr:s &&r) __name__ __module__ __qualname____doc__ __constants__int__annotations__r rrstrr __classcell__)rs@rrr sI DJM Kst5V55 'C'r) torch.nn.functionalnn functionalrtorchrmoduler__all__rrrr(s'  4'V4'r