add rocm extra

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "setuptools>=75.0.0,<76",
     "numpy>=1.26.3,<2",
     "pandas~=2.2.0",
-    "torch==2.5.1; sys_platform == 'darwin'",
+    "torch>=2.5.1",
     "pyarrow>=16.0.0",
     "joblib>=1.3.0,<2",
     "psutil>=5.9.5,<6",  # upgrade when colab psutil is updated
@@ -45,13 +45,19 @@ dependencies = [
 
 [project.optional-dependencies]
 gpu = [
-    "torch==2.5.1",
+    "torch>=2.5.1",
     "bitsandbytes>=0.43.3,<0.44",
     "vllm>=0.6.4,<0.7",
 ]
+rocm = [
+    "pytorch-triton-rocm>=3.1.0",
+    "torch>=2.5.1",
+    "torchvision>=0.20.1",
+    # "bitsandbytes>=0.43.3,<0.44",  # need to be built from source
+    # "vllm>=0.6.4,<0.7",  # need to be built from source
+]
 cpu = [
-    "torch==2.5.1+cpu; sys_platform != 'darwin'",
-    "torch==2.5.1; sys_platform == 'darwin'"
+    "torch>=2.5.1"
 ]
 
 [dependency-groups]
@@ -83,6 +89,7 @@ conflicts = [
   [
     { extra = "cpu" },
     { extra = "gpu" },
+    { extra = "rocm" },
   ],
 ]
 
@@ -91,10 +98,23 @@ name = "pytorch-cpu"
 url = "https://download.pytorch.org/whl/cpu"
 explicit = true
 
+[[tool.uv.index]]
+name = "pytorch-rocm"
+url = "https://download.pytorch.org/whl/rocm6.1"
+explicit = true
+
 [tool.uv.sources]
 torch = [
-  { index = "pytorch-cpu", extra = "cpu", marker = "sys_platform != 'darwin'" },
+  { index = "pytorch-cpu", extra = "cpu" },
+  { index = "pytorch-rocm", extra = "rocm" },
 ]
+torchvision = [
+  { index = "pytorch-rocm", extra = "rocm" },
+]
+pytorch-triton-rocm = [
+  { index = "pytorch-rocm", extra = "rocm" },
+]
+
 
 [tool.hatch.build.targets.sdist]
 include = ["mostlyai/engine"]