chore: upgrade deps + uv

Author: walln

.tool-versions

@@ -1 +1 @@
-rye 0.35.0
+uv  latest

pyproject.toml

@@ -7,57 +7,61 @@ authors = [
     { name = "Nick Wall", email = "46641379+walln@users.noreply.github.com" },
 ]
 dependencies = [
-    "flax>=0.8.4",
-    "torch>=2.2.2",
+    "flax>=0.9.1",
+    "torch>=2.5.1",
     "loguru>=0.7.2",
-    "matplotlib>=3.8.4",
-    "scikit-learn>=1.4.1.post1",
+    "matplotlib>=3.9.2",
+    "scikit-learn>=1.5.2",
     "numpy>=1.26.4",
-    "optax>=0.2.2",
-    "tqdm>=4.66.2",
-    "equinox>=0.11.3",
+    "optax>=0.2.5",
+    "tqdm>=4.66.6",
+    "equinox>=0.11.9",
     "orbax>=0.1.9",
-    "wandb>=0.17.2",
-    "jax-metrics>=0.2.4",
-    "typer>=0.12.2",
-    "jaxlib @ https://storage.googleapis.com/jax-releases/cuda12/jaxlib-0.4.26+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl",
-    "jax[cuda12_pip]>=0.4.26",
+    "wandb>=0.18.7",
+    "jax-metrics>=0.2.5",
+    "typer>=0.15.1",
+    "jax>=0.4.35",
+    "jaxlib>=0.4.35",
     "jmp>=0.0.4",
-    "datasets>=3.0.1",
-    "torchvision>=0.18.1",
-    "transformers>=4.42.3",
+    "datasets>=3.1.0",
+    "torchvision>=0.20.1",
+    "transformers>=4.46.3",
 ]
 license = "MIT"
 readme = "README.md"
 requires-python = ">= 3.12"
 
+[dependency-groups]
+cuda = ["jax[cuda12]>=0.4.35"]
+metal = ["jax-metal>=0.1.0"]
+cpu = ["jax[cpu]>=0.4.35"]
+dev = [
+    "pytest>=8.3.4",
+    "pre-commit>=4.0.1",
+    "nvitop>=1.3.2",
+    "ipykernel>=6.29.5",
+    "jupyter>=1.1.1",
+    "ruff>=0.8.4",
+]
+
+[tool.uv]
+default-groups = ["dev", "cpu"]
+conflicts = [[{ group = "cuda" }, { group = "metal" }]]
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
-[tool.uv]
-managed = true
-dev-dependencies = [
-    "pytest>=8.1.1",
-    "pre-commit>=3.7.0",
-    "nvitop>=1.3.2",
-    "ipykernel>=6.29.4",
-    "jupyter>=1.0.0",
-    "ruff>=0.4.10",
-]
-
 [tool.hatch.metadata]
 allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/scratch"]
 
 [tool.ruff]
+target-version = "py312"
 include = ["src/**", "tests/**"]
-
-[tool.ruff.lint]
-pydocstyle.convention = "google"
+line-length = 88
 select = [
     "E",      # pycodestyle
     "W",      # pycodestyle
@@ -67,17 +71,18 @@ select = [
     "D",      # pydocstyle
     "I",      # isort
     "SIM",    # flake8-simplify
-    "TCH",    # flake8-type-checking
     "TID",    # flake8-tidy-imports
     "UP",     # pyupgrade
     "PT",     # flake8-pytest-style
     "RUF",    # Ruff-specific rules
     "FBT001", # flake8-boolean-trap
 ]
-
 ignore = ["F722"]
 
-[tool.ruff.lint.isort]
+[tool.ruff.pydocstyle]
+convention = "google"
+
+[tool.ruff.isort]
 known-first-party = ["src"]
 
 [tool.pytest.ini_options]
@@ -86,8 +91,3 @@ filterwarnings = [
     'ignore:co_lnotab is deprecated, use co_lines instead.',
     'ignore:To copy construct from a tensor, it is recommended to use sourceTensor\.clone\(\)\.detach\(\) or sourceTensor\.clone\(\)\.detach\(\)\.requires_grad_\(True\):UserWarning',
 ]
-
-
-[tool.mypy]
-ignore_missing_imports = true
-strict = false