[MRG] Tensorflow backend & Benchmarker & Myst_parser (#316)

* First batch of tf methods (to be continued)

* Second batch of method (yet to debug)

* tensorflow for cpu

* add tf requirement

* pep8 + bug

* small changes

* attempt to solve pymanopt bug with tf2

* attempt #2

* attempt #3

* attempt 4

* docstring

* correct pep8 violation introduced in merge conflicts resolution

* attempt 5

* attempt 6

* just a random try

* Revert "just a random try"

This reverts commit 8223e768bfe33635549fb66cca2267514a60ebbf.

* GPU tests for tensorflow

* pep8

* attempt to solve issue with m2r2

* Remove transpose backend method

* first draft of benchmarker (need to correct time measurement)

* prettier bench table

* Bitsize and prettier device methods

* prettified table bench

* Bug corrected (results were mixed up in the final table)

* Better perf counter (for GPU support)

* pep8

* EMD bench

* solve bug if no GPU available

* pep8

* warning about tensorflow numpy api being required in the backend.py docstring

* Bug solve in backend docstring

* not covering code which requires a GPU

* Tensorflow gradients manipulation tested

* Number of warmup runs is now customizable

* typo

* Remove some warnings while building docs

* Change prettier_device to device_type in backend

* Correct JAX mistakes preventing to see the CPU if a GPU is present

* Attempt to solve JAX bug in case no GPU is found

* Reworked benchmarks order and results storage & clear GPU after usage by benchmark

* Add bench to backend docstring

* better benchs

* remove useless stuff

* Better device_type

* Now using MYST_PARSER and solving links issue in the README.md / online docs

1 files changed, 43 insertions, 1 deletions

diff --git a/test/test_gromov.py b/test/test_gromov.py
index 38a7fd7..4b995d5 100644
--- a/test/test_gromov.py
+++ b/test/test_gromov.py
@@ -9,7 +9,7 @@
 import numpy as np

 import ot

 from ot.backend import NumpyBackend

-from ot.backend import torch

+from ot.backend import torch, tf

 

 import pytest

 

@@ -113,6 +113,45 @@ def test_gromov_dtype_device(nx):
         nx.assert_same_dtype_device(C1b, gw_valb)

 

 

+@pytest.mark.skipif(not tf, reason="tf not installed")

+def test_gromov_device_tf():

+    nx = ot.backend.TensorflowBackend()

+    n_samples = 50  # nb samples

+    mu_s = np.array([0, 0])

+    cov_s = np.array([[1, 0], [0, 1]])

+    xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s, random_state=4)

+    xt = xs[::-1].copy()

+    p = ot.unif(n_samples)

+    q = ot.unif(n_samples)

+    C1 = ot.dist(xs, xs)

+    C2 = ot.dist(xt, xt)

+    C1 /= C1.max()

+    C2 /= C2.max()

+

+    # Check that everything stays on the CPU

+    with tf.device("/CPU:0"):

+        C1b = nx.from_numpy(C1)

+        C2b = nx.from_numpy(C2)

+        pb = nx.from_numpy(p)

+        qb = nx.from_numpy(q)

+        Gb = ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', verbose=True)

+        gw_valb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', log=False)

+        nx.assert_same_dtype_device(C1b, Gb)

+        nx.assert_same_dtype_device(C1b, gw_valb)

+

+    if len(tf.config.list_physical_devices('GPU')) > 0:

+        # Check that everything happens on the GPU

+        C1b = nx.from_numpy(C1)

+        C2b = nx.from_numpy(C2)

+        pb = nx.from_numpy(p)

+        qb = nx.from_numpy(q)

+        Gb = ot.gromov.gromov_wasserstein(C1b, C2b, pb, qb, 'square_loss', verbose=True)

+        gw_valb = ot.gromov.gromov_wasserstein2(C1b, C2b, pb, qb, 'kl_loss', log=False)

+        nx.assert_same_dtype_device(C1b, Gb)

+        nx.assert_same_dtype_device(C1b, gw_valb)

+        assert nx.dtype_device(Gb)[1].startswith("GPU")

+

+

 def test_gromov2_gradients():

     n_samples = 50  # nb samples

 

@@ -150,6 +189,7 @@ def test_gromov2_gradients():
 

 

 @pytest.skip_backend("jax", reason="test very slow with jax backend")

+@pytest.skip_backend("tf", reason="test very slow with tf backend")

 def test_entropic_gromov(nx):

     n_samples = 50  # nb samples

 

@@ -208,6 +248,7 @@ def test_entropic_gromov(nx):
 

 

 @pytest.skip_backend("jax", reason="test very slow with jax backend")

+@pytest.skip_backend("tf", reason="test very slow with tf backend")

 def test_entropic_gromov_dtype_device(nx):

     # setup

     n_samples = 50  # nb samples

@@ -306,6 +347,7 @@ def test_pointwise_gromov(nx):
     np.testing.assert_allclose(float(logb['gw_dist_std']), 0.0015952535464736394, atol=1e-8)

 

 

+@pytest.skip_backend("tf", reason="test very slow with tf backend")

 @pytest.skip_backend("jax", reason="test very slow with jax backend")

 def test_sampled_gromov(nx):

     n_samples = 50  # nb samples