Code

Add the JAX backend#504

Code

Open

YannCabanes wants to merge 14 commits intotslearn-team:maintslearn-team/tslearn:mainfrom

YannCabanes:add-jaxYannCabanes/tslearn:add-jaxCopy head branch name to clipboard

Commits

Commits on Dec 22, 2023

• Add the JAX backend

  YannCabanes

  committed
  0574fe5

  Copy full SHA for 0574fe5

Commits on Jan 9, 2024

• Merge branch 'main' into add-jax

  YannCabanes

  committed
  55dec26

  Copy full SHA for 55dec26

Commits on Jan 10, 2024

• Add the jax backend in backend.py and define jax random functions in jax_backend.py

  YannCabanes

  committed
  73cc1d1

  Copy full SHA for 73cc1d1

• Add tests for the jax backend

  YannCabanes

  committed
  5259437

  Copy full SHA for 5259437

• Add blank line

  YannCabanes

  committed
  fec935f

  Copy full SHA for fec935f

• Fix error in function select_backend when pytorch or jax is not installed

  YannCabanes

  committed
  f0ebbd0

  Copy full SHA for f0ebbd0

• Skit test_soft_dtw_loss_pytorch if torch is not installed

  YannCabanes

  committed
  214c18c

  Copy full SHA for 214c18c

Commits on Feb 5, 2024

• Improve JAX backend

  YannCabanes

  committed
  0555ee0

  Copy full SHA for 0555ee0

Commits on Feb 6, 2024

• Change input D of function _njit_soft_dtw_grad

  YannCabanes

  committed
  f1fd65a

  Copy full SHA for f1fd65a

Commits on Feb 23, 2024

• Fix index error in function _soft_dtw_grad

  YannCabanes

  committed
  a424f9b

  Copy full SHA for a424f9b

• Add test for pytorch and jax installation

  YannCabanes

  committed
  a4a56fa

  Copy full SHA for a4a56fa

• Add installation of the jax backend for tests in the file setup.py

  YannCabanes

  committed
  983e532

  Copy full SHA for 983e532

Commits on Feb 27, 2024

• Remove unnecessary zeros assignation in matrix D in _soft_dtw_grad

  YannCabanes

  committed
  ac78d86

  Copy full SHA for ac78d86

Commits on Apr 23, 2024

• Empty commit

  YannCabanes

  committed
  54f7ceb

  Copy full SHA for 54f7ceb