[MRG] Tensorflow backend & Benchmarker & Myst_parser

[ncassereau]

Personal To do

• Prettify benchmark results presentation
• Include benchmark for EMD
• Correct time measure for each backend (does time.perf_counter work on GPUs ??)
• Make real benchmark for Sinkhorn Knopp that is pretty for this PR (with every single backend and device)
• Put aforementioned benchmark in backend.py's docstring
• Check if set_gradient works as intended for tensorflow and add test
• Add to backend.py docstrings a comment about tensorflow requiring the numpy API
• Replace m2r2 by myst_parser (https://myst-parser.readthedocs.io/en/latest/)

Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Motivation and context / Related issue

• Backend for tensorflow.
• Also adds a benchmark folder which allows to easily benchmark every backend and present results in a html table (compatible with github comments and PR text).
• Also replaced m2r2 by myst_parser. It allows to use markdown files in our official docs. Myst_parser seems more maintained. It also happens to solve the issue with internal links introduced by the m2r2 update few PRs ago.

Note that tensorflow does not know how to transpose inplace, which can have a negative impact on performance for any function which requires the usage of a transpose (More information here: https://www.tensorflow.org/api_docs/python/tf/transpose#expandable-1)

Note that nothing was done regarding the device consistency of our calculation (unlike TorchBackend for instance). This is due to the fact that tensorflow chooses the device automatically without caring about the device where inputs are located. If the user wants a specific device, this should be done outside of POT, probably something like

with tf.device("/GPU:1"):
    ot.sinkhorn(args, **kwargs)

More information here: https://www.tensorflow.org/guide/gpu#overview

How has this been tested (if it applies)

Tested with Tensorflow 2.6.0 and 2.7.0. Benchmark with Tensorflow 2.6.0

Performance on Sinkhorn Knopp

The speedup is not as convincing as with cupy, might be due to how inefficient tensorflow transpose are. Still a great choice for very high dimensions.
Used pretty much the code from test/test_gpu.py::test_gpu_sinkhorn.
Tests were performed with Tensorflow 2.6.0, results in seconds, averaged over 100 runs.

Sinkhorn Knopp - Averaged on 100 runs
Bitsize	32 bits
Device	CPU	GPU
Sample size	Numpy	Pytorch	Tensorflow	Cupy	Jax	Pytorch	Tensorflow
50	0.0008	0.0022	0.0151	0.0095	0.0193	0.0051	0.0293
100	0.0005	0.0013	0.0097	0.0057	0.0115	0.0029	0.0173
500	0.0009	0.0016	0.0110	0.0058	0.0115	0.0029	0.0166
1000	0.0021	0.0021	0.0145	0.0056	0.0118	0.0029	0.0168
2000	0.0069	0.0043	0.0278	0.0059	0.0118	0.0030	0.0165
5000	0.0707	0.0314	0.1395	0.0074	0.0125	0.0035	0.0198
Bitsize	64 bits
Device	CPU	GPU
Sample size	Numpy	Pytorch	Tensorflow	Cupy	Jax	Pytorch	Tensorflow
50	0.0008	0.0020	0.0154	0.0093	0.0191	0.0051	0.0328
100	0.0005	0.0013	0.0094	0.0056	0.0114	0.0029	0.0169
500	0.0013	0.0017	0.0120	0.0059	0.0116	0.0029	0.0168
1000	0.0034	0.0027	0.0177	0.0058	0.0118	0.0029	0.0167
2000	0.0146	0.0075	0.0436	0.0059	0.0120	0.0029	0.0165
5000	0.1467	0.0568	0.2468	0.0077	0.0146	0.0045	0.0204

Code used for the test

python3 -m benchmarks.sinkhorn_knopp

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document.
• All tests passed, and additional code has been covered with new tests.

[codecov]

Codecov Report

Merging #316 (86faaa4) into master (b3dc68f) will decrease coverage by 0.31%.
The diff coverage is 88.75%.

@@            Coverage Diff             @@
##           master     #316      +/-   ##
==========================================
- Coverage   93.34%   93.03%   -0.32%     
==========================================
  Files          21       21              
  Lines        4901     5198     +297     
==========================================
+ Hits         4575     4836     +261     
- Misses        326      362      +36     

[rflamary]

This is awesome, thank you so much @ncassereau-idris