fix: jax v0.4.3.2 change to lax.convert_element_type_p

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 89.82%. Comparing base (93a833f) to head (a2ef913).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #143 +/- ## ========================================== - Coverage 89.83% 89.82% -0.01% ========================================== Files 25 25 Lines 1367 1366 -1 ========================================== - Hits 1228 1227 -1 Misses 139 139 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.