(1) Make config hashable. (2) Flax
This commit is contained in:
+1
-2
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user