# Vmap and its compatibility issues **URL:** \ **Category:** Uncategorized\ **Created:** [June 1, 2023, 4:58am UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169 "2023-06-01T04:58:23Z")\ **Posts on this page:** 5\ **Page:** 1
**Author:** ![hmdolatabadi](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/3ec8ea/32.png) [@hmdolatabadi](https://discuss.pytorch.org/u/hmdolatabadi)\ **Post date:** [June 1, 2023, 4:58am UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169/1 "2023-06-01T04:58:24Z")
Hi All, I have been using the `Python Optimal Transport` library. I want to define a loss function that iterates over every sample in my batch and calculates the `sinkhorn` distance for that sample and its ground-truth value. What I was using before was a for-loop: ```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) ``` but this is way too slow for my application. I was reading through `functorch`, and apparently I should have been able to use the `vmap` functionality. But after wrapping my function in `vmap`, I get this weird error that everyone else is talking about: ```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. ``` Does anyone have a workaround? Cheers. ---
**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\ **Post date:** [June 1, 2023, 7:51am UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169/2 "2023-06-01T07:51:14Z")
Maybe you could pull out the if condition and execute it before executing the loop to remove the data-dependency. ---
**Author:** ![AlphaBetaGamma96](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/32.png) [@AlphaBetaGamma96](https://discuss.pytorch.org/u/AlphaBetaGamma96)\ **Post date:** [June 1, 2023, 9:58am UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169/3 "2023-06-01T09:58:44Z")
Hi @hmdolatabadi, Following on from what @ptrblck said, `vmap` 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. What you should do is vmap over the `ot.sinkhorn2` function, then return a batch of outputs over which you sum (rather than summing within vmap itself). You can try this and see if it works, ```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) ``` FYI, if your `ot.sinkhorn2` function is an `nn.Module` object, you’ll need to use `torch.func.functional_call` with your vmap call. ---
**Author:** ![hmdolatabadi](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/3ec8ea/32.png) [@hmdolatabadi](https://discuss.pytorch.org/u/hmdolatabadi)\ **Post date:** [June 1, 2023, 10:48pm UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169/4 "2023-06-01T22:48:13Z")
> [@AlphaBetaGamma96](#): > > `in_dims=(0, 0, None, None)` Thanks @ptrblck and @AlphaBetaGamma96! The issue is more related to what @ptrblck described. The `sinkhorn` function in the `ot` package uses a data-dependent if-statement, and that’s why I’m getting an error: ```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) --> 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 . ``` I guess I might need an alternative route for computing the `sinkhorn` function that doesn’t involve if-statements. If you have any ideas, it is more than appreciated! Thanks. ---
**Author:** ![AlphaBetaGamma96](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/32.png) [@AlphaBetaGamma96](https://discuss.pytorch.org/u/AlphaBetaGamma96)\ **Post date:** [June 1, 2023, 11:14pm UTC](https://discuss.pytorch.org/t/vmap-and-its-compatibility-issues/181169/5 "2023-06-01T23:14:55Z")
If the `Python Optimal Transport` 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.