Consider switching to jax for the core functions

[kjappelbaum]

Feature description

We use numba to jit some of the core/utils functions right now. But we could just replace this dependency by jax. Additionally, we would get vectorization with vmap and potential parallelization with pmap. If we have a GPU, it will run all the matrix operations on the GPU.

Implementation idea

1. Replace import numpy as np with import jax.numpy as np and replace the numba jit decorator with jax jit compilation
2. Look into vectorizing some of the for loops

Alternatives

1. Stick to numba
2. Use cython for the localized for loops (pareto_classify)

[kjappelbaum]

I played a bit with it and it is a bit more effort than expected.

For example, we cannot just change the numpy import because we use inplace operations in several places in the code
Also the jit decorator threw a lot of errors

[github-actions]

Stale issue message