uv installation instructions

pyproject.toml

@@ -23,12 +23,12 @@ 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",
+    "flax",
+    "optax",
+    "einshape",
+    "orbax-checkpoint",
+    "jaxtyping",
+    "jax[cuda12]"
 ]
 
 [tool.ruff]

requirements.txt

@@ -5,13 +5,9 @@ certifi==2025.10.5
 charset-normalizer==3.4.3
     # via requests
 filelock==3.19.1
-    # via
-    #   huggingface-hub
-    #   torch
+    # via huggingface-hub
 fsspec==2025.9.0
-    # via
-    #   huggingface-hub
-    #   torch
+    # via huggingface-hub
 hf-xet==1.1.10
     # via huggingface-hub
 huggingface-hub==0.35.3
@@ -20,53 +16,8 @@ idna==3.10
     # via requests
 inquirerpy==0.3.4
     # via huggingface-hub
-jinja2==3.1.6
-    # via torch
-markupsafe==3.0.3
-    # via jinja2
-mpmath==1.3.0
-    # via sympy
-networkx==3.5
-    # via torch
 numpy==2.3.3
     # via timesfm (pyproject.toml)
-nvidia-cublas-cu12==12.8.4.1
-    # via
-    #   nvidia-cudnn-cu12
-    #   nvidia-cusolver-cu12
-    #   torch
-nvidia-cuda-cupti-cu12==12.8.90
-    # via torch
-nvidia-cuda-nvrtc-cu12==12.8.93
-    # via torch
-nvidia-cuda-runtime-cu12==12.8.90
-    # via torch
-nvidia-cudnn-cu12==9.10.2.21
-    # via torch
-nvidia-cufft-cu12==11.3.3.83
-    # via torch
-nvidia-cufile-cu12==1.13.1.3
-    # via torch
-nvidia-curand-cu12==10.3.9.90
-    # via torch
-nvidia-cusolver-cu12==11.7.3.90
-    # via torch
-nvidia-cusparse-cu12==12.5.8.93
-    # via
-    #   nvidia-cusolver-cu12
-    #   torch
-nvidia-cusparselt-cu12==0.7.1
-    # via torch
-nvidia-nccl-cu12==2.27.3
-    # via torch
-nvidia-nvjitlink-cu12==12.8.93
-    # via
-    #   nvidia-cufft-cu12
-    #   nvidia-cusolver-cu12
-    #   nvidia-cusparse-cu12
-    #   torch
-nvidia-nvtx-cu12==12.8.90
-    # via torch
 packaging==25.0
     # via huggingface-hub
 pfzy==0.3.4
@@ -79,20 +30,10 @@ requests==2.32.5
     # via huggingface-hub
 safetensors==0.6.2
     # via timesfm (pyproject.toml)
-setuptools==80.9.0
-    # via triton
-sympy==1.14.0
-    # via torch
-torch==2.8.0
-    # via timesfm (pyproject.toml)
 tqdm==4.67.1
     # via huggingface-hub
-triton==3.4.0
-    # via torch
 typing-extensions==4.15.0
-    # via
-    #   huggingface-hub
-    #   torch
+    # via huggingface-hub
 urllib3==2.5.0
     # via requests
 wcwidth==0.2.14