chore(deps): bump JAX version and use new jax.jit syntax

[jeertmans]

xref: jax-ml/jax#34697

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.74%. Comparing base (bd40505) to head (0de2252).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #370      +/-   ##
==========================================
- Coverage   84.77%   84.74%   -0.03%     
==========================================
  Files          32       32              
  Lines        3021     3016       -5     
==========================================
- Hits         2561     2556       -5     
  Misses        460      460              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[codspeed-hq]

CodSpeed Performance Report

Merging this PR will improve performance by ×32

_{Comparing fix-jax-jit (0de2252) with main (bd40505)}

Summary

⚡ 4 improved benchmarks
✅ 6 untouched benchmarks

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
⚡	test_compute_paths[exhaustive]	45.2 ms	40.6 ms	+11.21%
⚡	test_image_method	6,458.3 µs	203.7 µs	×32
⚡	test_compute_paths_disconnect_inactive_triangles_benchmark[True]	29.9 ms	24.5 ms	+22.23%
⚡	test_compute_paths_disconnect_inactive_triangles_benchmark[False]	44.3 ms	39.9 ms	+11.2%