Vectorizing random matrix creation using `vmap` with tensor of seeds

RS-Coop: <p>I have a situation where I want to apply a batch of random matrices to a batch of vectors, and I have a given batch of seeds that are used to construct each matrix. I do this by using <code>torch.manual_seed()</code> before generating the matrix. This seems ripe for <code>torch.vmap</code>, as it lets me vectorize over the batch dimension. However, the problem is that <code>torch.manual_seed()</code> converts the given seed to an integer, which implicitly calls <code>.item()</code> on a singleton tensor resulting in an error.</p>
<p>Here is some sample code implementing that idea.</p>
<pre data-code-wrap="python"><code class="lang-python">def _sketch(self, f, s):
    n = f.shape[0]
    r = round(0.01*n)

    torch.manual_seed(s)
    sketch = torch.randn(n, r, device='cpu')

    return torch.einsum('nc,nr-&gt;rc', f, sketch)

def sketch(self, features, seeds):
    return torch.vmap(self._sketch, randomness="different")(features, seeds)
</code></pre>
<p>It doesn’t seem to me like there should be any reason this can’t work. Any help here would be much appreciated!</p>

Soumya_Kundu: <p>Assuming this is the error you got:</p>
<p><code>&gt; RuntimeError: vmap: It looks like you're calling .item() on a Tensor. We don't support vmap over calling .item() on a Tensor, please try to rewrite what you're doing with other operations. If error is occurring somewhere inside PyTorch internals, please file a bug report.</code></p>
<p>I think internally, the .item() essentially is generating a scalar. So that is messing with the vmap I guess.</p>
<p>Replacing <code>torch.manual_seed(s)</code> with something of your own may work?</p>

AlphaBetaGamma96: <p>The <code>.item()</code> method converts a torch <code>Tensor</code> that’s a scalar to a <code>float</code>. Because <code>torch.func.vmap</code> vectorizes over pytorch primitives, if you leave the <code>torch</code> namespace (e.g. by using <code>.item()</code>), <code>torch.func.vmap</code> can no longer vectorize your function.</p>
<p>If you want to have different random samples with <code>torch.func.vmap</code> the <code>randomness="different"</code> kwarg will do that. I assume you’re manually selecting the seed for reproducibility?</p>
<p>In the docs it states you can get around this in 2 ways</p>
<ol>
<li>Generate all your random data outside of vmap, and just vectorize over the data</li>
<li>Define your own generator directly</li>
</ol>
<p>The docs on randomness are here: <a href="https://pytorch.org/docs/stable/func.ux_limitations.html#randomness" class="inline-onebox" rel="noopener nofollow ugc">UX Limitations — PyTorch 2.4 documentation</a></p>

RS-Coop: <p>Yes, that is the error. You are correct that <code>.item()</code> is called internally when converting the singleton integer tensor to a Python integer, which is causing said error.</p>
<p>Is there an obvious replacement for <code>manual_seed()</code> that still allows generating reproducible random tensors?</p>

RS-Coop: <p>Yes, the seed was used previously, elsewhere in my program, to generate a random matrix, and now I want to use the same seed to reproduce it. So the options of <code>randomness="different"</code> or <code>randomness="same"</code> don’t work since I want specific randomness.</p>
<p>I would like to avoid generating all the data outside of vmap, and that only solves half the problem since it still requires looping over the seeds to generate the batch of matrices.</p>
<p>I could define a generator object, but I still have to manually seed it if I want reproducibility. The seeds are also batched themselves from a larger set, so there is no distinct order, eliminating the possibility of seeding once outside of <code>vmap</code> and then relying on the reproducible sequence.</p>

Soumya_Kundu: <aside class="quote no-group" data-username="RS-Coop" data-post="4" data-topic="206977">
<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/rs-coop/48/71703_2.png" class="avatar"> RS-Coop:</div>
<blockquote>
<p>Is there an obvious replacement for <code>manual_seed()</code> that still allows generating reproducible random tensors?</p>
</blockquote>
</aside>
<p>Not that I know of <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"> A quick search really resulted in nothing. Even here: <a href="https://pytorch.org/docs/stable/notes/randomness.html" class="inline-onebox" rel="noopener nofollow ugc">Reproducibility — PyTorch 2.4 documentation</a></p>