fix JAX version <= 0.4.30
Fixes JAX version to <= 0.4.30, as newer versions severely degrade the performance of some demos, and some versions even produce segfaults in the pipeline.
Fixes JAX version to <= 0.4.30, as newer versions severely degrade the performance of some demos, and some versions even produce segfaults in the pipeline.