pulled flax changes from master

2025-10-07 21:46:30 +00:00
parent 018b2557ab fe70095c6c
commit 56e7fd67c6
11 changed files with 1276 additions and 19 deletions
@@ -16,8 +16,20 @@ dependencies = [
    "numpy>=1.26.4",
    "huggingface_hub[cli]>=0.23.0",
    "safetensors>=0.5.3",
+]
+
+[project.optional-dependencies]
+torch = [
    "torch>=2.0.0",
 ]
+flax = [
+    "flax>=0.8.2",
+    "optax>=0.2.2",
+    "einshape>=0.8.0",
+    "orbax-checkpoint>=0.5.15",
+    "jaxtyping>=0.2.29",
+    "jax>=0.4.26",
+]

 [tool.ruff]
 line-length = 88