Vmap and its compatibility issues

hmdolatabadi: <p>Hi All,</p>
<p>I have been using the <code>Python Optimal Transport</code> library. I want to define a loss function that iterates over every sample in my batch and calculates the <code>sinkhorn</code> distance for that sample and its ground-truth value. What I was using before was a for-loop:</p>
<pre><code class="lang-auto">for i in range(len(P_batch)):
      if i == 0:
         loss = ot.sinkhorn2(P_batch[i].view(-1, 1), Q_batch[i].view(-1, 1), C, epsilon)
      loss += ot.sinkhorn2(P_batch[i].view(-1, 1), Q_batch[i].view(-1, 1), C, epsilon)
</code></pre>
<p>but this is way too slow for my application. I was reading through <code>functorch</code>, and apparently I should have been able to use the <code>vmap</code> functionality. But after wrapping my function in <code>vmap</code>, I get this weird error that everyone else is talking about:</p>
<pre><code class="lang-auto">RuntimeError: vmap: It looks like you're attempting to use a Tensor in some data-dependent control flow. We don't support that yet, please shout over at https://github.com/pytorch/functorch/issues/257.
</code></pre>
<p>Does anyone have a workaround?</p>
<p>Cheers.</p>

ptrblck: <p>Maybe you could pull out the if condition and execute it before executing the loop to remove the data-dependency.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/hmdolatabadi">@hmdolatabadi</a>,</p>
<p>Following on from what <a class="mention" href="/u/ptrblck">@ptrblck</a> said, <code>vmap</code> will have a problem with the if statement, and it seems you’re only using the if statement to initialize the loss from the first sample, then in-place add all other samples in your batch.</p>
<p>What you should do is vmap over the <code>ot.sinkhorn2</code> function, then return a batch of outputs over which you sum (rather than summing within vmap itself).</p>
<p>You can try this and see if it works,</p>
<pre><code class="lang-auto">from torch.func import vmap

P = P_batch.flatten(start_dim=1)   #flatten tensors (but not including batch dim)
Q = Q_ batch.flatten(start_dim=1)

losses = vmap(ot.sinkhorn2, in_dims=(0, 0, None, None))(P, Q, C, epsilon) #assumes C and epsilon are constants (that are repeated for all samples). 
loss = torch.sum(losses)
</code></pre>
<p>FYI, if your <code>ot.sinkhorn2</code> function is an <code>nn.Module</code> object, you’ll need to use <code>torch.func.functional_call</code> with your vmap call.</p>

hmdolatabadi: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="3" data-topic="181169">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p><code>in_dims=(0, 0, None, None)</code></p>
</blockquote>
</aside>
<p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>!<br>
The issue is more related to what <a class="mention" href="/u/ptrblck">@ptrblck</a> described.<br>
The <code>sinkhorn</code> function in the <code>ot</code> package uses a data-dependent if-statement, and that’s why I’m getting an error:</p>
<pre><code class="lang-auto">File /anaconda3/envs/my_env/lib/python3.8/site-packages/ot/bregman.py:505, in sinkhorn_knopp(a, b, M, reg, numItermax, stopThr, verbose, log, warn, warmstart, **kwargs)
    502 v = b / KtransposeU
    503 u = 1. / nx.dot(Kp, v)
--&gt; 505 if (nx.any(KtransposeU == 0)
    506         or nx.any(nx.isnan(u)) or nx.any(nx.isnan(v))
    507         or nx.any(nx.isinf(u)) or nx.any(nx.isinf(v))):
    508     # we have reached the machine precision
    509     # come back to previous solution and quit loop
    510     warnings.warn('Warning: numerical errors at iteration %d' % ii)
    511     u = uprev

RuntimeError: vmap: It looks like you're attempting to use a Tensor in some data-dependent control flow. We don't support that yet, please shout over at https://github.com/pytorch/functorch/issues/257 .
</code></pre>
<p>I guess I might need an alternative route for computing the <code>sinkhorn</code> function that doesn’t involve if-statements. If you have any ideas, it is more than appreciated!</p>
<p>Thanks.</p>

AlphaBetaGamma96: <p>If the <code>Python Optimal Transport</code> library has a repo, I’d open an issue and ask, because it seems like you need a version that doesn’t have an if statement, unfortunately.</p>