Avoiding data copy when using array indexing

username: yiftach
cooked: <p>In the following toy example:</p>
<pre><code class="lang-auto">import torch
data = torch.randn(10, 1024)
indices = torch.tensor([0, 1, 7, 0])
selected_data = data[indices]
</code></pre>
<p>Is there a way to avoid a copy by having <code>selected_data</code> be a view of <code>data</code> and share the same underlying storage?</p>
<p>In my real case, as in this toy example, the indices in <code>indices</code> may have repeating and non following elements. However, I assume each element of <code>data</code> is large and <code>indices</code> contains many repeats.</p>
<p>I’ve seen other possible choices like <code>index_select()</code> but all seem to copy data. Can this be avoided?</p>
<p>Bonus: my <code>indices</code> actually have a more specific structure, with k&lt;n, it is [[0, 1 , …, n], [k, k + 1, …, n+k], [2k, 2k + 1, …, n+2k], …]. Does this somehow help? I assume that for [0, 1 , …, n] only it would, but the complete structure is more exotic.</p>
<p>Thanks!</p>

username: KFrank
cooked: <p>Hi Yiftach!</p>
<aside class="quote no-group" data-username="yiftach" data-post="1" data-topic="182748" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yiftach/48/58678_2.png" class="avatar"> yiftach:</div>
<blockquote>
<p>Is there a way to avoid a copy by having <code>selected_data</code> be a view of <code>data</code> and share the same underlying storage?</p>
</blockquote>
</aside>
<p>No, not with general index values.</p>
<p>Pytorch tensors requires that their elements be stored contiguously (although<br>
you can “stride” through the contiguous elements), so when you index into a<br>
tensor you have to create a new contiguous tensor, rather than an indexed<br>
view into the original tensor.</p>
<p>You can use <code>.untyped_storage().data_ptr()</code> to verify that new memory<br>
is being used.</p>
<aside class="quote no-group">
<blockquote>
<p>my <code>indices</code> actually have a more specific structure, with k&lt;n, it is [[0, 1 , …, n], [k, k + 1, …, n+k], [2k, 2k + 1, …, n+2k], …]. Does this somehow help?</p>
</blockquote>
</aside>
<p>For this specific use case your indices have a structure that makes it possible<br>
to create a strided view into your original tensor, rather than use explicit<br>
indexing.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; data = torch.tensor (list (range (10)))
&gt;&gt;&gt; indices = torch.tensor ([0, 1, 7, 0])
&gt;&gt;&gt; selected_data = data[indices]
&gt;&gt;&gt;
&gt;&gt;&gt; data.untyped_storage().data_ptr()                  # original storage for data
1811627391296
&gt;&gt;&gt; selected_data.untyped_storage().data_ptr()         # new storage for indexed data
1811621701056
&gt;&gt;&gt;
&gt;&gt;&gt; n = 4
&gt;&gt;&gt; k = 2
&gt;&gt;&gt;
&gt;&gt;&gt; indices = torch.tensor ([
...     [0, 1, 2, 3, 4],
...     [2, 3, 4, 5, 6],
...     [4, 5, 6, 7, 8]
... ])
&gt;&gt;&gt;
&gt;&gt;&gt; rows = (10 - (n + 1)) // k + 1
&gt;&gt;&gt; cols = n + 1
&gt;&gt;&gt;
&gt;&gt;&gt; dataA = data[indices]
&gt;&gt;&gt; dataB = data.as_strided ((rows, cols), (k, 1))
&gt;&gt;&gt;
&gt;&gt;&gt; data.untyped_storage().data_ptr()                  # original storage for data
1811627391296
&gt;&gt;&gt; dataA.untyped_storage().data_ptr()                 # index -- new storage
1811578654848
&gt;&gt;&gt; dataB.untyped_storage().data_ptr()                 # as_strided() -- view into original storage
1811627391296
&gt;&gt;&gt;
&gt;&gt;&gt; data[4] = 99                                       # modify original data
&gt;&gt;&gt; dataA                                              # doesn't change new tensor
tensor([[0, 1, 2, 3, 4],
        [2, 3, 4, 5, 6],
        [4, 5, 6, 7, 8]])
&gt;&gt;&gt; dataB                                              # change is reflected in view of original tensor
tensor([[ 0,  1,  2,  3, 99],
        [ 2,  3, 99,  5,  6],
        [99,  5,  6,  7,  8]])
</code></pre>
<p>Note that you have to be cautious when using such a view.  Quoting from the<br>
<a href="https://pytorch.org/docs/2.0/generated/torch.as_strided.html" rel="noopener nofollow ugc">as_strided() documentation:</a></p>
<blockquote>
<p>Warning</p>
<p>Prefer using other view functions, like <a href="https://pytorch.org/docs/2.0/generated/torch.Tensor.expand.html#torch.Tensor.expand" rel="noopener nofollow ugc"><code>torch.Tensor.expand()</code></a>, to setting a view’s strides manually with as_strided, as this function’s behavior depends on the implementation of a tensor’s storage. 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.</p>
</blockquote>
<p>In the example you give, your view is overlapped.  My <em>guess</em> is that if you<br>
treat <code>dataB</code> as read-only, you will be okay, but if you assign into <code>dataB</code>, you<br>
may get non-deterministic results for the values of the overlapped elements.</p>
<p>Best.</p>
<p>K. Frank</p>

username: yiftach
cooked: <p>That perfectly answers all of my concerns. Thanks!</p>

username: yiftach
cooked: <p>As a followup for the bonus case, what if my data is not one dimensional but n-dimensional? I’m working on videos where each one is of shape (channels, length, width, height) and normally I would do video[:, indices] but playing with as_strided doesn’t give me what I want - I’m modifying the size parameter to add C before and W, H after, and not sure what to add in the stride parameter:</p>
<pre><code class="lang-auto">import torch
data = torch.arange(10).repeat(2, 4, 4, 1).permute(0, 3, 1, 2)
indices = torch.tensor ([
     [0, 1, 2, 3, 4],
     [2, 3, 4, 5, 6],
     [4, 5, 6, 7, 8]
 ])
dataA = data[:, indices]

n = 4
k = 2
rows = (10 - (n + 1)) // k + 1
cols = n + 1
dataB = data.as_strided ((2, rows, cols, 4, 4), (1, k, 1, 1, 1))

dataB - dataA # not an array of zeros
</code></pre>
<p>So I wanted to make sure the option to share the view is still possible in this case.I would assume it is, but haven’t gotten it to work yet.</p>

username: KFrank
cooked: <p>Hi Yiftach!</p>
<aside class="quote no-group" data-username="yiftach" data-post="4" data-topic="182748" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/yiftach/48/58678_2.png" class="avatar"> yiftach:</div>
<blockquote>
<p>what if my data is not one dimensional but n-dimensional?</p>
</blockquote>
</aside>
<p>This will depend on the details of where your data comes from.</p>
<p>With <code>.as_strided()</code>, the <code>stride</code> you use is relative to the underlying<br>
storage of the tensor you call it on, rather than relative to the tensor<br>
itself (that is, how the tensor displays), which, itself, may have a<br>
non-trivial <code>stride</code>.</p>
<p>(I am using “trivial” or “default” <code>stride</code> to mean what pytorch calls a<br>
“contiguous” tensor.)</p>
<p>In any event, the <code>stride</code> you pass to <code>.as_strided</code> isn’t correct.  However,<br>
the correct <code>stride</code> will depend on what <code>stride</code> <code>data</code> has.</p>
<p>In the example you posted, <code>data</code>, being the result of a call to <code>.permute()</code>,<br>
does have a non-trivial stride, so that affects the <code>stride</code> you need to use<br>
to get your desired result.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; data = torch.arange(10).repeat(2, 4, 4, 1).permute(0, 3, 1, 2)
&gt;&gt;&gt; indices = torch.tensor ([
...      [0, 1, 2, 3, 4],
...      [2, 3, 4, 5, 6],
...      [4, 5, 6, 7, 8]
...  ])
&gt;&gt;&gt;
&gt;&gt;&gt; data._is_view()                # data is a view (into the contiguous tensor returned by .repeat())
True
&gt;&gt;&gt; data.is_contiguous()           # and so is not contiguous
False
&gt;&gt;&gt; data.stride()                  # and has "non-default" strides
(160, 1, 40, 10)
&gt;&gt;&gt;
&gt;&gt;&gt; dataA = data[:, indices]       # (the new tensor returned by indexing is contiguous)
&gt;&gt;&gt;
&gt;&gt;&gt; rows = 3
&gt;&gt;&gt; cols = 5
&gt;&gt;&gt; dataB = data.as_strided ((2, rows, cols, 4, 4), (1, 2, 1, 1, 1))           # the strides are wrong
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (dataB, dataA)     # so not equal
False
&gt;&gt;&gt;
&gt;&gt;&gt; dataCont = data.contiguous()   # make a contiguous version of data
&gt;&gt;&gt;
&gt;&gt;&gt; dataCont.is_contiguous()       # indeed contiguous
True
&gt;&gt;&gt; dataCont.stride()              # with "default" strides
(160, 16, 4, 1)
&gt;&gt;&gt;
&gt;&gt;&gt; dataC = dataCont.as_strided ((2, rows, cols, 4, 4), (160, 32, 16, 4, 1))   # strides are straightforward
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (dataC, dataA)     # works
True
&gt;&gt;&gt;
&gt;&gt;&gt; dataD = data.as_strided ((2, rows, cols, 4, 4), (160, 2, 1, 40, 10))       # strides are a little trickier
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (dataD, dataA)     # also works
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>