feature request: JAX JIT backend branch

[aidancrilly]

Very large scattering matrices are memory intensive so another approach might be to calculate scattering matrix elements when the spectrum is fully computed. This is more computationally expensive if repeatedly called so it is wise to leverage JIT compilation. It should be possible to write a python JAX version of the NeSST backend for this task. Need to figure out how to split the code up if this effort is pursued.