jax to 0.4.29; deprecate jax-md; dataset generation scripts; docs

[arturtoshev]

Upgrade JAX to the latest 0.4.28

[codecov[bot]]

Codecov Report

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

Project coverage is 82.1%. Comparing base (79b2db9) to head (fedaa4f).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #30 +/- ## ======================================= - Coverage 83.4% 82.1% -1.4% ======================================= Files 26 25 -1 Lines 1875 1680 -195 ======================================= - Hits 1565 1380 -185 + Misses 310 300 -10 ``` | [Flag](https://app.codecov.io/gh/tumaer/lagrangebench/pull/30/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tumaer) | Coverage Δ | | |---|---|---| | [unittests](https://app.codecov.io/gh/tumaer/lagrangebench/pull/30/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tumaer) | `82.1% <100.0%> (-1.4%)` | :arrow_down: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tumaer#carryforward-flags-in-the-pull-request-comment) to find out more.

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