Return `jax.Array`

[philip-paul-mueller]

JAX always returns jax.Array objects but we either return NumPy arrays or CuPy (although not implemented).
This will break drop in replacement.

If we are at it, we also have to modify the type annotation of the jit() wrapper.
Currently the return type is preserved, however, since we change that type we should put it to Any.

[philip-paul-mueller]

Solved by PR#26.