[MRG] Fix warning bug in sinkhorn2 (#417)

* Pass warn argument downstream in sinkhorn2 method.

* releases.md

* Fix unittest.

Co-authored-by: Rémi Flamary <[email protected]>

Author: arincbulgur

RELEASES.md

@@ -30,8 +30,9 @@ roughly 2^31) (PR #381)
 - Fixed an issue where the doc could not be built due to some changes in matplotlib's API (Issue #403, PR #402)
 - Replaced Numpy C Compiler with Setuptools C Compiler due to deprecation issues (Issue #408, PR #409)
 - Fixed weak optimal transport docstring (Issue #404, PR #410)
-- Fixed error whith parameter `log=True`for `SinkhornLpl1Transport` (Issue #412, 
+- Fixed error with parameter `log=True`for `SinkhornLpl1Transport` (Issue #412, 
 PR #413)
+- Fixed an issue about `warn` parameter in `sinkhorn2` (PR #417)
 - Fix an issue where the parameter `stopThr` in `empirical_sinkhorn_divergence` was rendered useless by subcalls 
   that explicitly specified `stopThr=1e-9` (Issue #421, PR #422).
 - Fixed a bug breaking an example where we would try to make an array of arrays of different shapes (Issue #424, PR #425)

ot/bregman.py

@@ -323,15 +323,18 @@ def sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000,
     if len(b.shape) < 2:
         if method.lower() == 'sinkhorn':
             res = sinkhorn_knopp(a, b, M, reg, numItermax=numItermax,
-                                 stopThr=stopThr, verbose=verbose, log=log,
+                                 stopThr=stopThr, verbose=verbose,
+                                 log=log, warn=warn,
                                  **kwargs)
         elif method.lower() == 'sinkhorn_log':
             res = sinkhorn_log(a, b, M, reg, numItermax=numItermax,
-                               stopThr=stopThr, verbose=verbose, log=log,
+                               stopThr=stopThr, verbose=verbose,
+                               log=log, warn=warn,
                                **kwargs)
         elif method.lower() == 'sinkhorn_stabilized':
             res = sinkhorn_stabilized(a, b, M, reg, numItermax=numItermax,
-                                      stopThr=stopThr, verbose=verbose, log=log,
+                                      stopThr=stopThr, verbose=verbose,
+                                      log=log, warn=warn,
                                       **kwargs)
         else:
             raise ValueError("Unknown method '%s'." % method)
@@ -344,15 +347,18 @@ def sinkhorn2(a, b, M, reg, method='sinkhorn', numItermax=1000,
 
         if method.lower() == 'sinkhorn':
             return sinkhorn_knopp(a, b, M, reg, numItermax=numItermax,
-                                  stopThr=stopThr, verbose=verbose, log=log,
+                                  stopThr=stopThr, verbose=verbose,
+                                  log=log, warn=warn,
                                   **kwargs)
         elif method.lower() == 'sinkhorn_log':
             return sinkhorn_log(a, b, M, reg, numItermax=numItermax,
-                                stopThr=stopThr, verbose=verbose, log=log,
+                                stopThr=stopThr, verbose=verbose,
+                                log=log, warn=warn,
                                 **kwargs)
         elif method.lower() == 'sinkhorn_stabilized':
             return sinkhorn_stabilized(a, b, M, reg, numItermax=numItermax,
-                                       stopThr=stopThr, verbose=verbose, log=log,
+                                       stopThr=stopThr, verbose=verbose,
+                                       log=log, warn=warn,
                                        **kwargs)
         else:
             raise ValueError("Unknown method '%s'." % method)

test/test_bregman.py

@@ -7,6 +7,7 @@
 #
 # License: MIT License
 
+import warnings
 from itertools import product
 
 import numpy as np
@@ -58,7 +59,10 @@ def test_convergence_warning(method):
         with pytest.warns(UserWarning):
             ot.barycenter(A, M, 1, method=method, stopThr=0, numItermax=1)
         with pytest.warns(UserWarning):
-            ot.sinkhorn2(a1, a2, M, 1, method=method, stopThr=0, numItermax=1)
+            ot.sinkhorn2(a1, a2, M, 1, method=method, stopThr=0, numItermax=1, warn=True)
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            ot.sinkhorn2(a1, a2, M, 1, method=method, stopThr=0, numItermax=1, warn=False)
 
 
 def test_not_implemented_method():