initial flax push

2025-10-03 14:14:36 -07:00
parent dc35e580d9
commit a314cd252c
10 changed files with 3526 additions and 4 deletions
@@ -14,11 +14,18 @@ readme = "README.md"
 packages = [{include = "timesfm", from = "src"}]

 [tool.poetry.dependencies]
-python = ">=3.11"
+python = ">=3.11,<4.0"
 numpy = ">=1.26.4"
 huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
 safetensors = ">=0.5.3"
-torch = { version = ">=2.0.0", extras = ["cuda"] }
+torch = { version = ">=2.0.0", extras = ["cuda"], optional = true }
+flax = { version = ">=0.8.2", optional = true }
+jax = { version = ">=0.4.26", optional = true }
+optax = { version = ">=0.2.2", optional = true }
+
+[tool.poetry.extras]
+torch = ["torch"]
+flax = ["flax", "jax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "typeguard"]

 [tool.ruff]
 line-length = 88