[MRG] Avoid changing precision in the backend (#572)

* Avoid changing precision in the backend

* Update RELEASES.md

Author: kachayev

RELEASES.md

@@ -22,6 +22,7 @@
 - Lazily instantiate backends to avoid unnecessary GPU memory pre-allocations on package import (Issue #516, PR #520)
 - Handle documentation and warnings when integers are provided to (f)gw solvers based on cg (Issue #530, PR #559)
 - Correct independence of `fgw_barycenters` to `init_C` and `init_X` (Issue #547, PR #566)
+- Avoid precision change when computing norm using PyTorch backend (Discussion #570, PR #572)
 
 ## 0.9.1
 *August 2023*

ot/backend.py

@@ -1941,7 +1941,7 @@ def power(self, a, exponents):
         return torch.pow(a, exponents)
 
     def norm(self, a, axis=None, keepdims=False):
-        return torch.linalg.norm(a.double(), dim=axis, keepdims=keepdims)
+        return torch.linalg.norm(a, dim=axis, keepdims=keepdims)
 
     def any(self, a):
         return torch.any(a)

test/test_gaussian.py

@@ -81,7 +81,7 @@ def test_empirical_bures_wasserstein_mapping_numerical_error_warning():
 
 
 def test_bures_wasserstein_distance(nx):
-    ms, mt = np.array([0]), np.array([10])
+    ms, mt = np.array([0]).astype(np.float32), np.array([10]).astype(np.float32)
     Cs, Ct = np.array([[1]]).astype(np.float32), np.array([[1]]).astype(np.float32)
     msb, mtb, Csb, Ctb = nx.from_numpy(ms, mt, Cs, Ct)
     Wb_log, log = ot.gaussian.bures_wasserstein_distance(msb, mtb, Csb, Ctb, log=True)