Is this kind of vectorization possible with vmap() or some torch function?

Ryan_Keathley1: <p>I am trying to simulate large numbers of 2d grid environments using GPU (a la Snake at a Million FPS).<br>
This involves frequently looking up a precomputed mask tensor at each step for each environment in the batch.</p>
<p>The environment is of shape (b, H, W) and the mask tensor is of shape (b, H, W, H, W) since we store a mask of the whole board for each square in each batch. Also this mask can be made to fit into memory by using a sparse tensor. since the masks are ‘small’.</p>
<p>So my first try was to perform a vmap on a tensor of indices <code>[[0, i_0, j_0], [1, i_1, j_1]...]</code><br>
using this function:</p>
<p><code>def get(t):     x=torch.select(masks, 0, t[0])     x=torch.select(x, 0, t[1])     x=torch.select(x, 0, t[2])     x=x.to_dense()     return x</code></p>
<p>which is essentially the sparse tensor version of <code>masks[t[0], t[1], t[2], :, :]</code></p>
<p>However this doesn’t work because vmap doesn’t support the use of <code>.item()</code>.<br>
This is disappointing because this tweet by pytorch has a similar use case with batch indexing</p>
<aside class="onebox twitterstatus" data-onebox-src="https://twitter.com/pytorch/status/1528800977108627456?lang=en">
  <header class="source">

      <a href="https://twitter.com/pytorch/status/1528800977108627456?lang=en" target="_blank" rel="noopener nofollow ugc">twitter.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://pbs.twimg.com/profile_images/1306686545974362113/JYq2LGIA_400x400.jpg" class="thumbnail onebox-avatar" width="400" height="400">
<h4><a href="https://twitter.com/pytorch/status/1528800977108627456?lang=en" target="_blank" rel="noopener nofollow ugc">PyTorch</a></h4>
<div class="twitter-screen-name"><a href="https://twitter.com/pytorch/status/1528800977108627456?lang=en" target="_blank" rel="noopener nofollow ugc">@PyTorch</a></div>

<div class="tweet">
  <span class="tweet-description">vmap lets you skip worrying about batch dimensions when writing code. For example, it’s generally tricky to batch code involving Tensor indexing. Now, don’t worry and just hand it off to vmap!
(3/n) <a href="https://t.co/vLYavPR2mk" target="_blank" rel="noopener nofollow ugc">https://t.co/vLYavPR2mk</a><div class="aspect-image-full-size" style="--aspect-ratio:666/500;"><img class="tweet-image" src="https://pbs.twimg.com/media/FTdkYWeUcAAWIFM.png" width="666" height="500"></div></span>
</div>

<div class="date">
  <a href="https://twitter.com/pytorch/status/1528800977108627456?lang=en" class="timestamp" target="_blank" rel="noopener nofollow ugc">6:12 PM - 23 May 2022</a>

    <span class="like">
      <svg viewbox="0 0 512 512" width="14px" height="16px" aria-hidden="true">
        <path d="M462.3 62.6C407.5 15.9 326 24.3 275.7 76.2L256 96.5l-19.7-20.3C186.1 24.3 104.5 15.9 49.7 62.6c-62.8 53.6-66.1 149.8-9.9 207.9l193.5 199.8c12.5 12.9 32.8 12.9 45.3 0l193.5-199.8c56.3-58.1 53-154.3-9.8-207.9z"></path>
      </svg>
      30
    </span>

    <span class="retweet">
      <svg viewbox="0 0 640 512" width="14px" height="16px" aria-hidden="true">
        <path d="M629.657 343.598L528.971 444.284c-9.373 9.372-24.568 9.372-33.941 0L394.343 343.598c-9.373-9.373-9.373-24.569 0-33.941l10.823-10.823c9.562-9.562 25.133-9.34 34.419.492L480 342.118V160H292.451a24.005 24.005 0 0 1-16.971-7.029l-16-16C244.361 121.851 255.069 96 276.451 96H520c13.255 0 24 10.745 24 24v222.118l40.416-42.792c9.285-9.831 24.856-10.054 34.419-.492l10.823 10.823c9.372 9.372 9.372 24.569-.001 33.941zm-265.138 15.431A23.999 23.999 0 0 0 347.548 352H160V169.881l40.416 42.792c9.286 9.831 24.856 10.054 34.419.491l10.822-10.822c9.373-9.373 9.373-24.569 0-33.941L144.971 67.716c-9.373-9.373-24.569-9.373-33.941 0L10.343 168.402c-9.373 9.373-9.373 24.569 0 33.941l10.822 10.822c9.562 9.562 25.133 9.34 34.419-.491L96 169.881V392c0 13.255 10.745 24 24 24h243.549c21.382 0 32.09-25.851 16.971-40.971l-16.001-16z"></path>
      </svg>
      3
    </span>
</div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


vmoens: <p>Thanks for this Ryan<br>
May I ask for some minimal reproducible example? It’ll be easier to see if we can come up with something (I’m pretty sure we can!)<br>
Few questions:</p>
<ul>
<li>Is it just the indexing of the mask that you want to vmap?</li>
<li>Also how big are b, H and W? Do you really need a sparse tensor for this, as values will be Boolean it should not take much space if b H and W are not crazy big.</li>
</ul>
<p>In the meantime I’ll try a few things given my understanding of the problem.</p>

Ryan_Keathley1: <p>This was my attempt and verification that its the indexing that’s the problem</p>
<p>import torch<br>
from functorch import vmap</p>
<p>idx = torch.zeros(size=(5, 0))<br>
values = torch.zeros(size=(0,))<br>
masks = torch.sparse_coo_tensor(idx,values,(10^4,20,20,20,20))</p>
<p>def get(t):<br>
x=torch.select(masks, 0, t[0])<br>
x=torch.select(x, 0, t[1])<br>
x=torch.select(x, 0, t[2])<br>
x=x.to_dense()<br>
return x</p>
<p>def get_(t):<br>
x=torch.select(masks, 0, 0)<br>
x=torch.select(x, 0, 1)<br>
x=torch.select(x, 0, 2)<br>
x=x.to_dense()<br>
return x</p>
<p>i = torch.tensor([[<em>, 0, 0] for _ in range(10^4)])<br>
y=vmap(get</em>)(i)<br>
print(y.shape)<br>
z=vmap(get)(i)<br>
print(z.shape)</p>
<p>There are other operations that must be batched, but they are more straightforward. I believe all the other rules of the environments can be implemented efficiently using convolutions and multiplication. These masks I’m retrieving are somewhat expensive to compute, so rather than doing that at every step I chose to memoize it. It should be quite large, if I’m to streamline this otherwise expensive RL algorithm (e.g. 10^5). Say H, W = 20, B = 10^4 for starters</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="Ryan_Keathley1" data-post="1" data-topic="162508">
<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/ryan_keathley1/48/53463_2.png" class="avatar"> Ryan_Keathley1:</div>
<blockquote>
<p>However this doesn’t work because vmap doesn’t support the use of <code>.item()</code>.</p>
</blockquote>
</aside>
<p>Open an issue on their github page and they might be able to fix it. Their github is <a href="https://github.com/pytorch/functorch/issues" rel="noopener nofollow ugc">here</a>.</p>

Ryan_Keathley1: <p>Did you manage to come up with anything? In the meantime, I’m exploring how to possibly circumvent this limitation. However it requires that the first n steps of my batch of states to all be the same, which is not ideal.</p>

vmoens: <p>I’m not sure what you’re trying to do can be done, but here are a few things we could explore:<br>
in your example you’re not using Boolean tensors. The shape of a Boolean tensor or 10<strong>4 x 20</strong>4 will be approx. 1.5Gb. It’s big, but not impossible to handle. Of course this is integrated is some more complex pipeline over which I have little grasp.</p>
<p>In TorchRL we have written a <a href="https://github.com/pytorch/rl/blob/main/torchrl/data/tensordict/memmap.py" rel="noopener nofollow ugc">MemmapTensor class</a> that you might be interested in. The idea is that you’re using a memory mapped array that sits on disk and that you can access easily in a row-by-row fashion (i.e. not loading the full data in memory to read it).<br>
To create it, just call <code>MemmapTensor(tensor)</code>.</p>
<p>Otherwise, we could perhaps work with the indices of the sparse tensor (match the indices from i with those from the sparse tensor)?</p>
<p>I’ll look in some other solutions too…</p>