[bug] Fixing compatibility issues with jax
(#691)
#1309
CI-models.yml
on: push
Matrix: test_linux
Matrix: test_macos
Matrix: test_windows