Add the JAX backend

[YannCabanes]

Add the JAX backend: https://jax.readthedocs.io/en/latest/ The JAX backend can be used for automatic differentiation: https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html

[YannCabanes]

The following error is obtained:

tslearn/tests/test_metrics.py:15: in <module>
    backends = [Backend("numpy"), None]
tslearn/backend/backend.py:99: in __init__
    self.backend = select_backend(data)
tslearn/backend/backend.py:75: in select_backend
    backends_instances = [NumPyBackend(), JAXBackend(), PyTorchBackend()]
tslearn/backend/backend.py:13: in __init__
    raise ValueError("Could not use JAX backend since JAX is not installed.")
E   ValueError: Could not use JAX backend since JAX is not installed.

[YannCabanes]

It seems from the previous error message that JAX is not installed during the continuous integration tests, therefore the tests in test_metrics.py are not running with the JAX backend.

On my local computer, these tests are failing with the following error message:

TypeError: '<class 'jaxlib.xla_extension.ArrayImpl'>' object does not support item assignment. JAX arrays are immutable. Instead of ``x[idx] = y``, use ``x = x.at[idx].set(y)`` or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html

[YannCabanes]

Difficulty: "JAX arrays are immutable" https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html

[YannCabanes]

A solution would be to create a class named JAXNumPyInterface which would define the operators of mutable objects. The Python operators are listed here: https://docs.python.org/3/library/operator.html