Implement Jax backend

In the simplest terms, to support Jax we have to

• add a @jit decorator like in Numba and
• optimistically replace import numpy as np with import jax.numpy as np

The API has of course some deviations, but for many cases this should work. Shall we add it here, alongside the Numba backend implementation?

• https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html#jax-vs-numpy
• https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html