Re-enable Jax 0.5.x

[rg936672]

PR Type

• Build related changes
• Tests

Description

Re-allowed jax 0.5.x, ran uv sync --upgrade, and fixed all failing tests.

Some errors were just due to precision issues, which I resolved by tweaking the allowed precision difference.

The other two errors have been marked with TODOs; these should be fixed fairly urgently as they seem to be genuine mathematical errors rather than just precision issues.

How Has This Been Tested?

All tests pass.

Does this PR introduce a breaking change?

No.

Checklist before requesting a review

• I have made sure that my PR is not a duplicate.
• My code follows the style guidelines of this project.
• I have ensured my code is easy to understand, including docstrings and comments where necessary.
• I have performed a self-review of my code.
• I have made corresponding changes to the documentation.
• My changes generate no new warnings.
• New and existing unit tests pass locally with my changes.
• Any dependent changes have been merged and published in downstream modules.
• I have updated CHANGELOG.md, if appropriate.

[github-actions]

Performance review

Commit e009087 - Merge dd3fac5 into 6ed8e9d

Statistically significant changes

• basic_herding:
  • OLD: compilation 0.7294 units ± 0.09183 units; execution 0.1997 units ± 0.009922 units
  • NEW: compilation 1.005 units ± 0.08205 units; execution 0.179 units ± 0.009369 units
  • Significant increase in compilation time (37.71%, p=1.477e-06)
  • Significant decrease in execution time (-10.37%, p=0.0001448)
• basic_rpc:
  • OLD: compilation 0.7972 units ± 0.03942 units; execution 0.1277 units ± 0.003236 units
  • NEW: compilation 1.183 units ± 0.04287 units; execution 0.1216 units ± 0.001873 units
  • Significant increase in compilation time (48.44%, p=4.918e-14)
• basic_stein:
  • OLD: compilation 5.754 units ± 0.3955 units; execution 0.5415 units ± 0.01056 units
  • NEW: compilation 6.256 units ± 0.323 units; execution 0.4212 units ± 0.01473 units
  • Significant decrease in execution time (-22.21%, p=3.112e-13)

---

Normalisation values for new data:
Compilation: 1 unit = 347.71 ms
Execution: 1 unit = 805.7 ms

[rg936672]

Stein thinning errors seem to be fixed by #1000.

[github-actions]

Performance review

Commit 84bcdde - Merge 4fd3875 into c7fe19d

Statistically significant changes

• basic_herding:
  • OLD: compilation 0.7257 units ± 0.1075 units; execution 0.2017 units ± 0.012 units
  • NEW: compilation 1.062 units ± 0.1199 units; execution 0.188 units ± 0.01434 units
  • Significant increase in compilation time (46.29%, p=3.604e-06)
• basic_rpc:
  • OLD: compilation 0.7924 units ± 0.03353 units; execution 0.1284 units ± 0.000796 units
  • NEW: compilation 1.233 units ± 0.04242 units; execution 0.1216 units ± 0.0005996 units
  • Significant increase in compilation time (55.61%, p=4.062e-15)
• basic_stein:
  • OLD: compilation 5.973 units ± 0.2911 units; execution 0.5546 units ± 0.01259 units
  • NEW: compilation 6.654 units ± 0.3641 units; execution 0.4262 units ± 0.01417 units
  • Significant increase in compilation time (11.41%, p=0.0002374)
  • Significant decrease in execution time (-23.14%, p=3.948e-14)

---

Normalisation values for new data:
Compilation: 1 unit = 342.07 ms
Execution: 1 unit = 813.74 ms

[tp832944]

@rg936672 approved subject to merge conflicts.

[github-actions]

Performance review

Commit 4ebd6c5 - Merge 41ac74b into 25e32fc

Statistically significant changes

• basic_herding:
  • OLD: compilation 0.6896 units ± 0.06902 units; execution 0.1925 units ± 0.007103 units
  • NEW: compilation 0.9823 units ± 0.06289 units; execution 0.183 units ± 0.01087 units
  • Significant increase in compilation time (42.46%, p=1.104e-08)
• basic_rpc:
  • OLD: compilation 0.7939 units ± 0.1032 units; execution 0.1266 units ± 0.0009522 units
  • NEW: compilation 1.187 units ± 0.03916 units; execution 0.02285 units ± 0.000307 units
  • Significant increase in compilation time (49.56%, p=1.403e-07)
  • Significant decrease in execution time (-81.94%, p=4.971e-23)
• basic_stein:
  • OLD: compilation 5.689 units ± 0.4153 units; execution 0.5392 units ± 0.01237 units
  • NEW: compilation 5.485 units ± 0.4138 units; execution 0.4273 units ± 0.01774 units
  • Significant decrease in execution time (-20.76%, p=1.896e-11)

---

Normalisation values for new data:
Compilation: 1 unit = 334.57 ms
Execution: 1 unit = 800.35 ms