make jax optional install?

[nkern]

as of a few months ago an import jax statement was added to abscal.py without any exceptions, making jax a requirement for building hera_cal. I don't think jax is particularly hard to build, but it is a somewhat niche dependency and not needed in general to run most of the useful utilities in hera_cal. can we make this an optional dependency with a try: except: statement whereever jax is imported?

alternatively, if ppl feel that jax is an indispensable and widely used tool within hera_cal (which I don't think is the case), then let's formally make jax a dependency in the docs.

[jsdillon]

We've been lax in updating the docs, but I do believe it's in the configs for when you pip install. I think we want to move toward using jax more as we find more places it allows speed-ups, so I'm generally in favor of making it mandatory (it's a lot less heavy and finicky to install than something like tensorflow or pytorch).