[MRG] Test Jax version in backend (#794)

* Fix test jax_version in backend

* Fix test jax_version in backend

Author: clbonet

RELEASES.md

@@ -21,6 +21,7 @@ This new release adds support for sparse cost matrices and a new lazy EMD solver
 - Fix openmp flags on macOS (PR #789)
 - Clean documentation (PR #787)
 - Fix code coverage (PR #791)
+- Fix test of the version of jax in `ot.backend` (PR #794)
 
 
 ## 0.9.6.post1

ot/backend.py

@@ -122,7 +122,12 @@
         from jax.extend.backend import get_backend as _jax_get_backend
 
         jax_type = jax.numpy.ndarray
-        jax_new_version = float(".".join(jax.__version__.split(".")[1:])) > 4.24
+        jax_new_version = tuple([float(s) for s in jax.__version__.split(".")]) > (
+            0,
+            4,
+            24,
+            0,
+        )
     except ImportError:
         jax = False
         jax_type = float