How to divide a tensor with size MxN to smaller tensors

username: isalirezag
cooked: <p>I am interested in dividing a tensor to subtensors.<br>
eg. if i have tensor <code> A = torch.rand(2,3,8,8)</code> i what to generate tenosr <code>B</code> which has the size of <code>2x3x16x2x2</code> and then convert it back to in a way that has the same size as <code>A</code>.</p>
<p>in  this case, the dimension of tensor A is divided by k = 4 and we chose subblock size of s=8//k=2 (2x2).</p>
<p>in other words, the subblock are chosen in the following way when the same color means those elements are belonging to the same subblock</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/4/4e397ffafa50ea42c0870b9dc1181b251a6ef09e.png" data-download-href="https://discuss.pytorch.org/uploads/default/4e397ffafa50ea42c0870b9dc1181b251a6ef09e" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/4/4e397ffafa50ea42c0870b9dc1181b251a6ef09e.png" alt="image" data-base62-sha1="ba0r0SwXP9tI5TOJbgEQRWAihDU" width="452" height="500" data-dominant-color="9B748C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">625×690 1.88 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>im not sure how to do it in a fast way, (the brutforce way is using looops), i think the fold unfold should be faster.</p>
<p>can anyone help me please how i can code it in a fast way? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

username: vmirly1
cooked: <p>You can reshape tensor <code>A</code>:</p>
<pre><code class="lang-python">&gt;&gt;&gt; A = torch.rand(2,3,8,8)
&gt;&gt;&gt; B = A.view(2, 3, 16, 2, 2)
&gt;&gt;&gt; B.shape
torch.Size([2, 3, 16, 2, 2])
</code></pre>

username: isalirezag
cooked: <p>This wont do the reshaping in a way that i want. the view function will reshape row wise.<br>
here is how it will be if we do the way that you say<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/1/1537ea1c1c6de688d2bc1f1d3e9af43c400f46b6.png" alt="image" data-base62-sha1="31HP25HjpN4K2vWxJe5YHuLOrNc" width="611" height="404"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/9/9ad4b868ae9dd3054c0bbbd17521d9118059cb66.png" data-download-href="https://discuss.pytorch.org/uploads/default/9ad4b868ae9dd3054c0bbbd17521d9118059cb66" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9ad4b868ae9dd3054c0bbbd17521d9118059cb66_2_458x500.png" alt="image" data-base62-sha1="m5HhsOWCjkpNNU5T6EMASLpL9TU" width="458" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9ad4b868ae9dd3054c0bbbd17521d9118059cb66_2_458x500.png, https://discuss.pytorch.org/uploads/default/original/2X/9/9ad4b868ae9dd3054c0bbbd17521d9118059cb66.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/9/9ad4b868ae9dd3054c0bbbd17521d9118059cb66.png 2x" data-dominant-color="FCFCFC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">625×681 23 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>but i need the first subblock to be like this</p>
<pre><code class="lang-auto">
[[0.803,0.6546],
[0.2339,0.9017]]
</code></pre>

username: vmirly1
cooked: <p>I have a solution but maybe not the most elegant one. First, let’s work with the following tensor for A to make it easier :</p>
<pre><code class="lang-python">&gt;&gt;&gt; A = torch.tensor(np.arange(2*3*64)/10.0).reshape(2, 3, 8, 8)
&gt;&gt;&gt; A.shape
torch.Size([2, 3, 8, 8])
&gt;&gt;&gt; A[0, 0]
tensor([[0.0000, 0.1000, 0.2000, 0.3000, 0.4000, 0.5000, 0.6000, 0.7000],
        [0.8000, 0.9000, 1.0000, 1.1000, 1.2000, 1.3000, 1.4000, 1.5000],
        [1.6000, 1.7000, 1.8000, 1.9000, 2.0000, 2.1000, 2.2000, 2.3000],
        [2.4000, 2.5000, 2.6000, 2.7000, 2.8000, 2.9000, 3.0000, 3.1000],
        [3.2000, 3.3000, 3.4000, 3.5000, 3.6000, 3.7000, 3.8000, 3.9000],
        [4.0000, 4.1000, 4.2000, 4.3000, 4.4000, 4.5000, 4.6000, 4.7000],
        [4.8000, 4.9000, 5.0000, 5.1000, 5.2000, 5.3000, 5.4000, 5.5000],
        [5.6000, 5.7000, 5.8000, 5.9000, 6.0000, 6.1000, 6.2000, 6.3000]],
       dtype=torch.float64)
</code></pre>
<p>and then we can take B as follows:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; temp = A.view(2, 3, 32, 2)
&gt;&gt;&gt; ind  = np.arange(32)
&gt;&gt;&gt; ind2 = ind.reshape(8, 4).T.reshape(16, 2)
&gt;&gt;&gt; B = temp[:, :, ind2, :]
&gt;&gt;&gt; B2.shape
torch.Size([2, 3, 16, 2, 2])


&gt;&gt;&gt; B[0, 0, 0]
tensor([[0.0000, 0.1000],
        [0.8000, 0.9000]], dtype=torch.float64)
</code></pre>
<p>So the blocks in <code>B</code> are created as intended. The next 2x2 block is</p>
<pre><code class="lang-python">&gt;&gt;&gt; B[0, 0, 1]
tensor([[1.6000, 1.7000],
        [2.4000, 2.5000]], dtype=torch.float64)
</code></pre>
<p>And for the 2x2 block in the second column of A:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; B2[0, 0, 4]
tensor([[0.2000, 0.3000],
        [1.0000, 1.1000]], dtype=torch.float64)
</code></pre>

username: isalirezag
cooked: <p>thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"> it is a little slow, but still faster than brute force way</p>

username: vmirly1
cooked: <p>No problem!</p>
<p>Try replacing the numpy arrays with torch and see if that makes a difference. Note that if you are doing this in a loop, you can just create the tensor for indices once and reuse that in the loop.</p>

username: darrenjkt
cooked: <p>Thank you for this solution!</p>
<p>I would like to know how I can extend this 2D block solution for a 3D block? I’ve been trying a couple reshaping sizes but can’t seem to get the same result that you did.</p>

username: johnsutor1
cooked: <p>Another potential solution, at least for image tensors</p>
<pre><code class="lang-python">def unstack(tensor: torch.Tensor, dimensions: int):
    """
    Unstacks the provided image tensor into the specified amount of dimensions. 
    Expects a tensor in the shape of bs x c x n x n, where bs is the batch size,
    c is the number of channels. Unpacks the image to the size 
    (n / dimensions) ^2 * bs, c, dimensions, dimensions

    Parameters:
    - tensor (torch.Tensor): Image tensor to unstack
    - dimensions (int): Dimensions to unpack to.

    Returns
    - tensor (torch.Tensor): Unstacked image tensor
    """

    assert tensor.shape[-2] == tensor.shape[-1], "Image tensor must be square" 
    assert tensor.shape[-1] % dimensions == 0, "Cannot unpack the image into that shape"

    # Unfold the tensor to the specified dimensions
    tensor = tensor.unfold(1, 1, 1).unfold(2,dimensions,dimensions).unfold(3,dimensions,dimensions)

    # Reshape the tensor to be the shape
    tensor = tensor.squeeze().permute(2,1,0,3,4).reshape(-1, 3, dimensions,dimensions)

    return tensor
</code></pre>