Minimal JAX version bump

[EiffL]

This PR aims for minimal edits allowing for jax version update. With this update, we can reach the latest jax 0.8 version, quadax prevents us from going to jax v0.9.

TFP is apparentently no longer releasing new versions, but tfp-nightly continues to be maintained for new jax releases. We should still ultimately remove TFP because it's only used for this very speicific Bessel function, and it prevents installs on python 3.13, but for now this would be an easy fix and we can then compare the code speed before and after changing the bessel function implementation on same version of jax.

I left the tfp-nighlty version unconstrained, it should automatically adjust to the jax version, which in turn will be the latest version permitted by quadrax.

[codspeed-hq]

CodSpeed Performance Report

Merging this PR will degrade performance by 35.47%

_{Comparing jax_version_update (dca40a7) with main (f1dcdd7)}

Summary

⚡ 5 improved benchmarks
❌ 10 (👁 9) regressed benchmarks
✅ 16 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
⚡	test_benchmark_interpimage_flux_frac[run]	2.7 ms	2 ms	+33.6%
👁	test_benchmarks_lanczos_interp[xval-no_conserve_dc-compile]	79.5 ms	100.6 ms	-20.93%
👁	test_benchmark_moffat_init[compile]	437.1 ms	677.4 ms	-35.47%
👁	test_benchmark_interpimage_flux_frac[compile]	343.8 ms	484 ms	-28.97%
❌	test_benchmarks_lanczos_interp[xval-conserve_dc-run]	1 ms	1.3 ms	-21.63%
👁	test_benchmarks_interpolated_image[compile]	1.7 s	2.5 s	-31.35%
👁	test_benchmarks_lanczos_interp[kval-no_conserve_dc-compile]	108.5 ms	137.6 ms	-21.13%
⚡	test_benchmark_rng_discard[run]	1,261.5 µs	564.9 µs	×2.2
👁	test_benchmark_rng_discard[compile]	170.6 ms	213.4 ms	-20.05%
⚡	test_benchmark_spergel_calcfluxrad[compile]	15.5 s	11.9 s	+30.66%
👁	test_benchmarks_lanczos_interp[kval-conserve_dc-compile]	108.5 ms	137.4 ms	-21.05%
⚡	test_benchmark_invert_ab_noraise[compile]	1.7 s	1.1 s	+47.68%
⚡	test_benchmarks_interpolated_image[run]	58 ms	47.6 ms	+21.8%
👁	test_benchmark_spergel_conv[compile]	2.4 s	3.6 s	-34.04%
👁	test_benchmark_spergel_conv[run]	1.8 s	2.6 s	-30.74%