Add Jax bindings

[ceriottm]

Tagging @felixmusil who seemed interested. Challenge here would be to reuse as much as possible the CUDA code @nickjbrowning wrote for libtorch. The whole concept is to keep the codebase as compact as possible.