Add the JAX backend

azure-pipelines.yml

@@ -30,6 +30,7 @@ jobs:
       python --version
       python -m pip install --upgrade pip
       python -m pip install -r requirements.txt
+      python -m pip install jax
       python -m pip install torch
     displayName: 'Install dependencies'
 
@@ -46,7 +47,7 @@ jobs:
       python -m pytest -v tslearn/ --doctest-modules
     displayName: 'Test'
 
-- job: 'linux_without_torch'
+- job: 'linux_without_torch_and_jax'
   pool:
     vmImage: 'ubuntu-latest'
   strategy:
@@ -69,6 +70,7 @@ jobs:
       python --version
       python -m pip install --upgrade pip
       python -m pip install -r requirements.txt
+      python -m pip uninstall jax
       python -m pip uninstall torch
     displayName: 'Install dependencies'
 
@@ -110,6 +112,7 @@ jobs:
       python --version
       python -m pip install --upgrade pip
       python -m pip install -r requirements.txt
+      python -m pip install jax
       python -m pip install torch
     displayName: 'Install dependencies'
 
@@ -163,6 +166,7 @@ jobs:
       export OPENBLAS=$(brew --prefix openblas)
       python -m pip install --upgrade pip
       python -m pip install -r requirements.txt
+      python -m pip install jax
       python -m pip install torch
     displayName: 'Install dependencies'
 
@@ -209,6 +213,7 @@ jobs:
       python --version
       python -m pip install --upgrade pip
       python -m pip install -r requirements_nocast.txt
+      python -m pip install jax
       python -m pip install torch
     displayName: 'Install dependencies'
 

setup.py

@@ -35,7 +35,7 @@
     packages=find_packages(),
     package_data={"tslearn": [".cached_datasets/Trace.npz"]},
     install_requires=['numpy', 'scipy', 'scikit-learn', 'numba', 'joblib'],
-    extras_require={'tests': ['pytest', 'torch'], 'pytorch': ['torch']},
+    extras_require={'tests': ['pytest', 'torch', 'jax'], 'pytorch': ['torch'], 'jax': ['jax']},
     version=VERSION,
     url="http://tslearn.readthedocs.io/",
     project_urls={

tslearn/backend/backend.py

@@ -2,6 +2,17 @@
 
 from tslearn.backend.numpy_backend import NumPyBackend
 
+try:
+    import jax
+
+    from tslearn.backend.jax_backend import JAXBackend
+except ImportError:
+
+    class JAXBackend:
+        def __init__(self):
+            raise ValueError("Could not use JAX backend since JAX is not installed.")
+
+
 try:
     import torch
 
@@ -11,7 +22,7 @@
     class PyTorchBackend:
         def __init__(self):
             raise ValueError(
-                "Could not use pytorch backend since torch is not installed"
+                "Could not use PyTorch backend since Torch is not installed."
             )
 
 
@@ -28,14 +39,14 @@ def instantiate_backend(*args):
     backend : Backend instance
         The backend instance.
     """
+    backends_str = ["numpy", "jax", "torch"]
     for arg in args:
         if isinstance(arg, Backend):
             return arg
-        arg_string = (str(type(arg)) + str(arg)).lower()
-        if "numpy" in arg_string:
-            return Backend("numpy")
-        if "torch" in arg_string:
-            return Backend("pytorch")
+        arg_str = (str(type(arg)) + str(arg)).lower()
+        for backend_str in backends_str:
+            if backend_str in arg_str:
+                return Backend(backend_str)
     return Backend("numpy")
 
 
@@ -54,11 +65,15 @@ def select_backend(data):
         The backend class.
         If data is a Numpy array or data equals 'numpy' or data is None,
         backend equals NumpyBackend().
-        If data is a PyTorch array or data equals 'pytorch',
+        If data is a JAX array or data equals 'jax',
+        backend equals JAXBackend().
+        If data is a PyTorch tensor or data equals 'pytorch',
         backend equals PytorchBackend().
     """
-    arg_string = (str(type(data)) + str(data)).lower()
-    if "torch" in arg_string:
+    arg_str = (str(type(data)) + str(data)).lower()
+    if "jax" in arg_str:
+        return JAXBackend()
+    if "torch" in arg_str:
         return PyTorchBackend()
     return NumPyBackend()
 
@@ -72,6 +87,8 @@ class Backend(object):
         Indicates the backend to choose.
         If data is a Numpy array or data equals 'numpy' or data is None,
         self.backend is set to NumpyBackend().
+        If data is a JAX array or data equals 'jax',
+        self.backend is set to JAXBackend().
         If data is a PyTorch array or data equals 'pytorch',
         self.backend is set to PytorchBackend().
         Optional, default equals None.
@@ -85,6 +102,7 @@ def __init__(self, data=None):
                 setattr(self, element, getattr(self.backend, element))
 
         self.is_numpy = self.backend_string == "numpy"
+        self.is_jax = self.backend_string == "jax"
         self.is_pytorch = self.backend_string == "pytorch"
 
     def get_backend(self):
@@ -103,20 +121,24 @@ def cast(data, array_type="numpy"):
         The input data should be a list or numpy array or torch array.
         The data to cast.
     array_type: string
-        The type to cast the data. It can be "numpy", "pytorch" or "list".
+        The type to cast the data. It can be "numpy", "jax", "pytorch" or "list".
 
     Returns
     --------
     data_cast: array-like
         Data cast to array_type.
     """
-    data_type_string = str(type(data)).lower()
+    data_type_str = str(type(data)).lower()
     array_type = array_type.lower()
     if array_type == "pytorch":
         array_type = "torch"
-    if array_type in data_type_string:
+    if array_type in data_type_str:
         return data
     if array_type == "list":
         return data.tolist()
     be = Backend(array_type)
+    backends_str = ["numpy", "jax", "torch"]
+    for backend_str in backends_str:
+        if backend_str in data_type_str:
+            data = data.tolist()
     return be.array(data)