Documentation bug and feature suggestion for ChannelShuffle

KFrank: <p>Hi Forum!</p>
<p>As I read the <a href="https://pytorch.org/docs/2.3/generated/torch.nn.ChannelShuffle.html#torch.nn.ChannelShuffle" rel="noopener nofollow ugc">ChannelShuffle documentation,</a> it implies that the <code>[-3]</code><br>
dimension (the “C” dimension) of the input tensor will be shuffled.  But<br>
this is not the observed behavior.  For example:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.3.1'
&gt;&gt;&gt; torch.nn.ChannelShuffle (2) (torch.ones (1, 2, 1, 1))      # agrees with doc
tensor([[[[1.]],

         [[1.]]]])
&gt;&gt;&gt; torch.nn.ChannelShuffle (2) (torch.ones (1, 2, 1, 1, 1))   # disagrees with doc
tensor([[[[[1.]]],


         [[[1.]]]]])
&gt;&gt;&gt; torch.nn.ChannelShuffle (2) (torch.ones (1, 2, 1))         # disagrees with doc
tensor([[[1.],
         [1.]]])
&gt;&gt;&gt; torch.nn.ChannelShuffle (2) (torch.ones (1, 1, 2, 1, 1))   # disagrees with doc
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "&lt;path_to_pytorch_install&gt;\torch\nn\modules\module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;\torch\nn\modules\module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "&lt;path_to_pytorch_install&gt;\torch\nn\modules\channelshuffle.py", line 54, in forward
    return F.channel_shuffle(input, self.groups)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Number of channels must be divisible by groups. Got 1 channels and 2 groups.
</code></pre>
<p>I’m not sure what behavior is intended, but either way, the documentation<br>
should match the behavior.</p>
<p>(As noted in <a href="https://github.com/pytorch/pytorch/issues/123053" rel="noopener nofollow ugc">github issue #123053,</a>, the example code in the documentation<br>
is incorrect (specifically, <code>randn()</code> doesn’t produce integer values)).</p>
<p>The feature suggestion is to make <code>ChannelShuffle</code> more flexible (and<br>
in my mind simpler).  Let <code>ChannelShuffle</code> operate on tensors with<br>
arbitrary numbers of dimensions, and add a <code>dim</code> argument that specifies<br>
which dimension to shuffle.  (One can shuffle any desired dimension by<br>
using the appropriate <code>unsqueeze()</code>s and <code>permute()</code>s, but this is less<br>
readable and a nuisance.)</p>
<p>If the new argument <code>dim</code> is given the default value <code>dim = -3</code>, the default<br>
behavior will be backward-compatible with the <em>documented</em> behavior,<br>
although not with the <em>observed</em> behavior.</p>
<p>(Sorry for not posting this as a github issue – google nuked the account<br>
I had been using to post on github.)</p>
<p>Best.</p>
<p>K. Frank</p>

soulitzer: <aside class="quote no-group" data-username="KFrank" data-post="1" data-topic="204711">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>the example code in the documentation<br>
is incorrect (specifically, <code>randn()</code> doesn’t produce integer values)</p>
</blockquote>
</aside>
<p>(btw this should have been fixed on main!)</p>