refactored training approaches

2026-05-31 16:43:36 +00:00 · 2026-02-19 18:23:08 +01:00
parent 5912062dc0
commit 1a9901f118
8 changed files with 947 additions and 308 deletions
--- a/engine/jax/requirements.txt
+++ b/engine/jax/requirements.txt
@@ -1,5 +1,5 @@
-flax>=0.8.0
-optax>=0.2.0
-distrax>=0.1.5
-orbax-checkpoint>=0.5.0
-chex>=0.1.8
+flax==0.10.7
+optax==0.2.7
+distrax==0.1.5
+orbax-checkpoint==0.11.32
+chex==0.1.90