Make tests compatible with newer version of JAX

[ricardoV94]

Tests started failing with JAX 0.4.1 in #121

[codecov-commenter]

Codecov Report

Merging #133 (78939ab) into main (16d1cbe) will increase coverage by 0.05%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #133      +/-   ##
==========================================
+ Coverage   74.41%   74.46%   +0.05%     
==========================================
  Files         179      179              
  Lines       49249    49292      +43     
  Branches    10422    10432      +10     
==========================================
+ Hits        36649    36706      +57     
+ Misses      10295    10282      -13     
+ Partials     2305     2304       -1     

Impacted Files	Coverage Δ
pytensor/sparse/basic.py	82.53% <ø> (ø)
pytensor/sparse/rewriting.py	75.51% <ø> (ø)
pytensor/graph/basic.py	88.13% <100.00%> (+0.03%)	⬆️
pytensor/link/jax/dispatch/random.py	100.00% <100.00%> (ø)
pytensor/link/jax/dispatch/scalar.py	97.50% <100.00%> (+0.09%)	⬆️
pytensor/tensor/basic.py	89.95% <100.00%> (ø)
pytensor/tensor/blas.py	79.67% <100.00%> (ø)
pytensor/tensor/io.py	78.15% <100.00%> (ø)
pytensor/tensor/math.py	90.40% <100.00%> (ø)
pytensor/tensor/shape.py	91.74% <100.00%> (ø)
... and 3 more

[michaelosthege]

Looks like this is going to fix the CI on main

[ricardoV94]

Ugh, no idea what happened with the commit author.