(1) Make config hashable. (2) Flax

This commit is contained in:
siriuz42
2025-10-06 05:58:30 +00:00
parent c159f1bf39
commit 25bb4f20e8
8 changed files with 15 additions and 20 deletions
+1 -2
View File
@@ -24,13 +24,12 @@ optax = { version = ">=0.2.2", optional = true }
einshape = { version = ">=0.8.0", optional = true }
orbax-checkpoint = { version = ">=0.5.15", optional = true }
jaxtyping = { version = ">=0.2.29", optional = true }
typeguard = { version = ">=4.3.0", optional = true }
jax = { version = ">=0.4.26", optional = true }
[tool.poetry.extras]
torch = ["torch"]
flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "typeguard", "jax"]
flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "jax"]
[tool.ruff]
line-length = 88