From 7e0ea27ad9cad31cfc2181430d837c0a77a61568 Mon Sep 17 00:00:00 2001
From: Clément Bonet <32179275+clbonet@users.noreply.github.com>
Date: Fri, 5 May 2023 10:53:48 +0200
Subject: [MRG] Fix bug SSW backend (#471)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix bug np vs torch matmul

* typo error

* einsum projections ssw

* Test broadcast matmul

* einsum projections ssw

* Test broadcast matmul

* projections SSW with einsum

* reduce number of samples in test wasserstein_circle_unif

* Update releases.md

---------

Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
---
 ot/backend.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

(limited to 'ot/backend.py')

diff --git a/ot/backend.py b/ot/backend.py
index eecf9dd..d661c74 100644
--- a/ot/backend.py
+++ b/ot/backend.py
@@ -959,6 +959,14 @@ class Backend():
         """
         raise NotImplementedError()
 
+    def matmul(self, a, b):
+        r"""
+        Matrix product of two arrays.
+
+        See: https://numpy.org/doc/stable/reference/generated/numpy.matmul.html#numpy.matmul
+        """
+        raise NotImplementedError()
+
 
 class NumpyBackend(Backend):
     """
@@ -1293,6 +1301,9 @@ class NumpyBackend(Backend):
             return args[0]
         return args
 
+    def matmul(self, a, b):
+        return np.matmul(a, b)
+
 
 class JaxBackend(Backend):
     """
@@ -1645,6 +1656,9 @@ class JaxBackend(Backend):
             return jax.lax.stop_gradient((args[0],))[0]
         return [jax.lax.stop_gradient((a,))[0] for a in args]
 
+    def matmul(self, a, b):
+        return jnp.matmul(a, b)
+
 
 class TorchBackend(Backend):
     """
@@ -2098,6 +2112,9 @@ class TorchBackend(Backend):
             return args[0].detach()
         return [a.detach() for a in args]
 
+    def matmul(self, a, b):
+        return torch.matmul(a, b)
+
 
 class CupyBackend(Backend):  # pragma: no cover
     """
@@ -2474,6 +2491,9 @@ class CupyBackend(Backend):  # pragma: no cover
             return args[0]
         return args
 
+    def matmul(self, a, b):
+        return cp.matmul(a, b)
+
 
 class TensorflowBackend(Backend):
 
@@ -2865,3 +2885,6 @@ class TensorflowBackend(Backend):
         if len(args) == 1:
             return tf.stop_gradient(args[0])
         return [tf.stop_gradient(a) for a in args]
+
+    def matmul(self, a, b):
+        return tnp.matmul(a, b)
-- 
cgit v1.2.3