fix(lib): make this library compatible with recent JAX versions

[jeertmans]

No description provided.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (59cdf8f) to head (709d816).
Report is 22 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main       #81   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            9         9           
  Lines          976       978    +2     
=========================================
+ Hits           976       978    +2     

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

[github-actions]

Benchmark results:

Benchmark	before	changes
accumulate_on_transmitters_grid_over_paths_noapprox	32.6 ms	33.9 ms: 1.04x slower
accumulate_on_transmitters_grid_over_paths_approx	37.7 ms	34.2 ms: 1.10x faster
path_method_image	42.7 us	336 us: 7.88x slower
path_method_FPT	77.3 us	377 us: 4.88x slower
path_method_MPT	108 us	413 us: 3.83x slower
Geometric mean	(ref)	2.68x slower