Make JAX Optional

[jeremykubica]

Make JAX an optional dependency.

We currently only use JAX in a few places, so installing it on every instance seems unnecessary. This PR makes JAX and optional dependency and checks that it can be imported in the places where it is used. There is some additional logic needed to handle the BasicMathNode as that needs to have the library available in the globals dictionary to use with the eval function.

[hombit]

Thank you!

At some point I'd like us to start testing all dependencies with lower-version dependency resolution, I believe it is not part of the PPT standard testing suite

[github-actions]

Before [0220525]	After [cc7eec8]	Ratio	Benchmark (Parameter)
127±5μs	110±10μs	~0.87	benchmarks.TimeSuite.time_evaluate_simple_linear_wavelength_model
12.4±0.1μs	12.6±0.2μs	1.02	benchmarks.TimeSuite.time_sample_x1_from_hostmass
11.0±0.04ms	11.1±0.1ms	1.01	benchmarks.TimeSuite.time_make_evaluate_constant_sed_model
566±10μs	563±10μs	1.00	benchmarks.TimeSuite.time_apply_passbands
19.4±0.1ms	19.5±1ms	1.00	benchmarks.TimeSuite.time_chained_evaluate_sed
4.74±0.05ms	4.75±0.04ms	1.00	benchmarks.TimeSuite.time_lightcurve_source
28.9±0.4μs	28.8±0.3μs	1.00	benchmarks.TimeSuite.time_make_simple_linear_wavelength_model
33.8±0.3ms	33.6±0.6ms	0.99	benchmarks.TimeSuite.time_additive_multi_model_source
7.58±0.05ms	7.52±0.04ms	0.99	benchmarks.TimeSuite.time_apply_white_noise
4.61±0.06ms	4.56±0.1ms	0.99	benchmarks.TimeSuite.time_evaluate_salt3_passbands

Click here to view all benchmarks.