Make JAX optional

[xiki-tempula]

Pymbar is making JAX optional with pymbar-core (no JAX) and pymbar (with JAX)
choderalab/pymbar#500

Do we want to have a similar arrangement of alchemlyb-core and alchemlyb?

[orbeckst]

We could, although until someone really needs the slimmer, slower *-core version, I'd say we use high-performance by default where possible.

[orbeckst]

Note issue #336 regarding JAX issues.