Efficient way to concat different size tensors without using a loop

rubcovvasilii: <p>Hi there,</p>
<p>I have a list of different size tensors, and I want to concat some of them using indices without using for-loop.<br>
Here is an example to do it using a loop:</p>
<pre><code class="lang-auto">item_features = [
    torch.LongTensor([0]),
    torch.LongTensor([1, 10]),
    torch.LongTensor([2, 27, 31]),
    torch.LongTensor([3]),
    torch.LongTensor([4]),
    torch.LongTensor([5, 11]),
]

item_ids = torch.LongTensor([1, 5])

result = torch.cat([
    item_features[item_id]
    for item_id in item_ids 
])  
result  # tensor([ 1, 10,  5, 11])
</code></pre>
<p>Also I tried to do it by storing the tensors as a sparse csr matrix:</p>
<pre><code class="lang-auto">item_features = torch.LongTensor([0, 1, 10, 2, 27, 31, 3, 4, 5, 11])
item_indices = torch.LongTensor([0, 1, 3, 6, 7, 8, 10])

item_ids = torch.LongTensor([1, 5])

result = torch.cat([
    item_features[item_indices[item_id]: item_indices[item_id + 1]]
    for item_id in item_ids 
])
result  # tensor([ 1, 10,  5, 11])
</code></pre>
<p>But there is a problem of multiple slicing, which is similar to this topic: <a href="https://discuss.pytorch.org/t/how-to-slice-multiple-interval-over-tensor/163562" class="inline-onebox">How to slice multiple interval over tensor</a>.</p>
<p>So, how can I do this in an efficient tensor-like way?</p>

KFrank: <p>Hi Vasiliy!</p>
<aside class="quote no-group" data-username="rubcovvasilii" data-post="1" data-topic="191864" 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/rubcovvasilii/48/65480_2.png" class="avatar"> rubcovvasilii:</div>
<blockquote>
<p>I have a list of different size tensors</p>
</blockquote>
</aside>
<p>If you start with a list of tensors, you will need to loop over that list one<br>
way or another.</p>
<aside class="quote no-group">
<blockquote>
<p>Also I tried to do it by storing the tensors as a sparse csr matrix:</p>
</blockquote>
</aside>
<p>Pytorch tensor indexing with a boolean index tensor returns a tensor<br>
whose shape is not that of the index tensor.  Instead it returns a<br>
flattened tensor whose length is the number of <code>True</code>s in the index<br>
tensor.</p>
<p>If you are willing to start with your csr format, boolean indexing will<br>
let you get around the fact that individual elements of <code>item_features</code><br>
differ in length, making it possible to avoid looping over <code>item_ids</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">result = torch.cat([
    item_features[item_indices[item_id]: item_indices[item_id + 1]]
    for item_id in item_ids 
])
</code></pre>
<p>So, how can I do this in an efficient tensor-like way?</p>
</blockquote>
</aside>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.1.0
&gt;&gt;&gt;
&gt;&gt;&gt; item_features = torch.LongTensor([0, 1, 10, 2, 27, 31, 3, 4, 5, 11])
&gt;&gt;&gt; item_indices = torch.LongTensor([0, 1, 3, 6, 7, 8, 10])
&gt;&gt;&gt;
&gt;&gt;&gt; item_ids = torch.LongTensor([1, 5])
&gt;&gt;&gt;
&gt;&gt;&gt; # version with loop (list comprehension)
&gt;&gt;&gt; result = torch.cat([
...     item_features[item_indices[item_id]: item_indices[item_id + 1]]
...     for item_id in item_ids
... ])
&gt;&gt;&gt; result  # tensor([ 1, 10,  5, 11])
tensor([ 1, 10,  5, 11])
&gt;&gt;&gt;
&gt;&gt;&gt; # loop-free version using boolean indexing
&gt;&gt;&gt; seq = torch.arange (len (item_features)).unsqueeze (0)
&gt;&gt;&gt; inds = torch.logical_and (seq &gt;= item_indices[:-1, None][item_ids], seq &lt; item_indices[1:,  None][item_ids])
&gt;&gt;&gt; resultB = item_features.expand (inds.shape)[inds]
&gt;&gt;&gt; resultB
tensor([ 1, 10,  5, 11])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (resultB, result)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

rubcovvasilii: <p>Hi, KFrank! Thank you for the answer!</p>
<blockquote>
<p>If you start with a list of tensors, you will need to loop over that list one<br>
way or another.</p>
</blockquote>
<p>Yes, I started with a list, which isn’t good structure, and I’m thinking about how to store my features in an efficient way.<br>
For the solution with torch.logical_and - it creates a tensor with len(item_features) x len(item_ids) size and this is a new challenge. During my training I have ~10_000_000 features and ~100_000 item_ids (which is batch size) and I can’t create 10_000_000 x 100_000 tensor.</p>

rubcovvasilii: <p>Seems that I found a solution</p>
<pre><code class="lang-auto">item_features = torch.LongTensor([0, 1, 10, 2, 27, 31, 3, 4, 5, 11])
# item_indices[item_id] is an index where item_id features started in item_features
item_indices = torch.LongTensor([0, 1, 3, 6, 7, 8])  
# item_features_sizes[item_id] is a number of item_id features
item_features_sizes = torch.LongTensor([1, 2, 3, 1, 1, 2])


item_ids = torch.LongTensor([1, 5])  # batch

indices = item_indices[item_ids]
sizes = item_features_sizes[item_ids]
prep_indices = indices - torch.cat([
    torch.LongTensor([0]), 
    sizes.cumsum(0)[:-1],
])
feature_indices = (
    torch.repeat_interleave(prep_indices, sizes) + torch.arange(sizes.sum()) 
)
result = item_features[feature_indices]  # tensor([ 1, 10,  5, 11])
</code></pre>
<p>But maybe there is a more elegant way to do that.</p>

KFrank: <p>Hi Vasiliy!</p>
<aside class="quote no-group" data-username="rubcovvasilii" data-post="4" data-topic="191864" 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/rubcovvasilii/48/65480_2.png" class="avatar"> rubcovvasilii:</div>
<blockquote>
<p>Seems that I found a solution</p>
</blockquote>
</aside>
<p>That looks sound.</p>
<aside class="quote no-group">
<blockquote>
<p>But maybe there is a more elegant way to do that.</p>
</blockquote>
</aside>
<p>I really don’t think there is – at least I don’t see one.</p>
<p>At issue is the fact that pytorch tensors don’t support* “ragged” tensors<br>
(that is tensors whose individual rows differ in length).  So one way or<br>
another you have to pluck differing numbers of elements out of<br>
<code>item_features</code>.</p>
<p>Your solution does that by using <code>repeat_interleave()</code>.  (My approach<br>
uses boolean indexing.)  There’s bound to be some “inelegant” processing<br>
in order to implement that step.</p>
<p>*)  Pytorch does have <a href="https://pytorch.org/docs/2.0/sparse.html" rel="noopener nofollow ugc">sparse</a> tensors, but support for them is pretty limited.<br>
Perhaps you could try that and hope that as you implemented your<br>
algorithm all of the tensor operations you needed were, in fact, supported<br>
for sparse tensors.</p>
<p>Best.</p>
<p>K. Frank</p>