How to apply vmap on a heterogeneous tensor

Roy-Kid: <p>Hi,<br>
I have a kernel that works well on non-batch data, and I want to extend it to batchwise-format, <strong>and calculate gradient through it</strong>. The problem is our batch is not heterogeneous, it looks like:</p>
<pre><code class="lang-auto">x = torch.tensor([1, 2, 2, 3, 3, 3.])
batch = torch.tensor([0, 1, 1, 2, 2, 2])
natoms = torch.bincount(batch)
stride = torch.tensor([0, 1, 3, 6])
</code></pre>
<p>what we want to do is use vmap instead of for loop:</p>
<pre><code class="lang-auto">result = []
for i in stride:
    ans = kernel(x[stride: stride+1])
    result.append(ans)
return torch.concat(result)

vmap(kernel, in_dims=(None, 0, 0))(x, stride[:-1], stride[1:])
</code></pre>
<p>apparently it’s not working. What should I do? is there any other solution except padding?</p>

soulitzer: <p>You might be interested in nested tensors which are special subclass of tensor designed to work on batches of data where the batches are jagged/have varying length.</p>
<pre><code class="lang-auto">torch.nested.nested_tensor_from_jagged(values=x, offsets=stride)

# NestedTensor(size=(3, j1), offsets=tensor([0, 1, 3, 6]), contiguous=True)
</code></pre>
<p>There’s a tutorial here:<br>
<a href="https://pytorch.org/tutorials/prototype/nestedtensor.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/prototype/nestedtensor.html</a></p>
<p>(Note that some parts of it are out of date, any call to <a href="https://pytorch.org/docs/stable/nested.html#torch.nested.nested_tensor" rel="noopener nofollow ugc">torch.nested.nested_tensor</a> should be passed a <code>layout=torch.jagged</code> argument; The part about implementing MHA  has been updated relatively recently though)</p>
<p>NestedTensors unfortunately don’t work with vmap yet, so you’ll need to explicitly write out the batch dim in your program. (We have plans to support this soon though)</p>

Roy-Kid: <p>Thank you so much! I’ve used nested tensors in my <code>collate_fn</code> before, but I haven’t tried applying it in this context yet. I’ll test it out and update this post once I have some results.</p>

Roy-Kid: <p>I have read the doc but am still slightly confused: when should I use jagged?</p>
<pre><code class="lang-auto">    # Note: the torch.jagged layout is a nested tensor layout that supports a single ragged
    # dimension and works with torch.compile. The batch items each have shape (B, S*, D)
    # where B = batch size, S* = ragged sequence length, and D = embedding dimension.
</code></pre>
<p>here is my typical atom coordinates:</p>
<pre><code class="lang-auto">nested_tensor([
  tensor([[0, 0, 0],
          [0, 0, 1],
          [0, 1, 0]]),
  tensor([[0, 0, 0],
          [0, 0, 1],
          [0, 1, 0],
          [1, 0, 0],
          [0, 1, 1]])
])
</code></pre>
<p>I believe those tensors I don’t need a jagged layout. But if I have a bunch of tensors with shape <code>[(n, ), (m, )...]</code>, then I need jagged?  Is that jagged mean the dimension in the middle is irregular?</p>

soulitzer: <blockquote>
<p>I believe those tensors I don’t need a jagged layout.</p>
</blockquote>
<p>Yep</p>
<blockquote>
<p>But if I have a bunch of tensors with shape <code>[(n, ), (m, )...]</code> , then I need jagged?</p>
</blockquote>
<p>Jagged layout only supports a single jagged dimension, e.g. among the tensors that you have, only a single dimension can vary. So if you have a bunch of 2-D tensors where n/m can both vary, then jagged layout would not work.</p>
<blockquote>
<p>Is that jagged mean the dimension in the middle is irregular?</p>
</blockquote>
<p>Your jagged dimension cannot be the very first one because there needs to be a preceeding batch dimension.<br>
But you can have a 2D nested tensor, e.g. with shape= [3, [1, 2, 3]] where the varying dimension is the very last one.</p>

Roy-Kid: <p><img src="https://discuss.pytorch.org/images/emoji/apple/rofl.png?v=12" title=":rofl:" class="emoji" alt=":rofl:" loading="lazy" width="20" height="20"> I find the tracker <a href="https://github.com/pytorch/pytorch/issues/112398" class="inline-onebox" rel="noopener nofollow ugc">[Tracker] Move nested tensors to beta · Issue #112398 · pytorch/pytorch · GitHub</a>. Is there any chance you could provide an estimate for when <code>vmap</code> will support nested tensors? I’d really love to see this feature as well.</p>