Does PyTorch support stride_tricks as in numpy.lib.stride_tricks.as_strided?

Question

It is possible to make cool things by changing the strides of an array in Numpy like this:

import numpy as np
from numpy.lib.stride_tricks import as_strided

a = np.arange(15).reshape(3,5)

print(a)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10 11 12 13 14]]

b = as_strided(a, shape=(3,3,3), strides=(a.strides[-1],)+a.strides)

print(b)
# [[[ 0  1  2]
#   [ 5  6  7]
#   [10 11 12]]

#  [[ 1  2  3]
#   [ 6  7  8]
#   [11 12 13]]

#  [[ 2  3  4]
#   [ 7  8  9]
#   [12 13 14]]]


# Get 3x3 sums of a, for example
print(b.sum(axis=(1,2)))
# [54 63 72]

I searched a similar method in PyTorch and found as_strided, but it does not support strides which makes an element have multiple indices referring to it, as the warning says:

The constructed view of the storage must only refer to elements within the storage or a runtime error will be thrown, and if the view is “overlapped” (with multiple indices referring to the same element in memory) its behavior is undefined.

In particular it says that the behavior is undefined for the example above where elements have multiple indices.

Is there a way to make this work with documented, specified behavior?

Answer 1

How about torch.nn.unfold Extracts sliding local blocks from a batched input tensor.

Documentation is here: https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html

And like so I have been able to replicate the behavior you have shown in your code:

    a = torch.arange(15).reshape(3, 5)

    print("Original tensor:")
    print(a)

    # Unfold the tensor to create a view with a specified size and step
    unfolded = a.unfold(1, 3, 1).unfold(0, 3, 1)

    print("\nUnfolded tensor:")
    print(unfolded)

    sums = unfolded.sum(dim=(2, 3))

    print("\n3x3 sums:")
    print(sums) # tensor([[54, 63, 72]])

Warning: Currently, only 4-D input tensors (batched image-like tensors) are supported.