aqlaboratory · jandom · Apr 28, 2026 · Mar 23, 2026 · Mar 23, 2026 · Mar 23, 2026
diff --git a/docs/imgs/pixi-environments-and-features.excalidraw b/docs/imgs/pixi-environments-and-features.excalidraw
diff --git a/docs/imgs/pixi-environments-and-features.png b/docs/imgs/pixi-environments-and-features.png
diff --git a/docs/source/Installation.md b/docs/source/Installation.md
@@ -31,6 +31,7 @@ pixi run -e openfold3-cpu run_openfold
 We provide the following environments:
  - openfold3-cpu (linux-64, linux-aarch64, osx-64,osx-arm64)
  - openfold3-cuda12 and openfold-cuda13 (linux-64, linux-aarch64) 
+ - openfold3-rocm7 (linux-64)
 
 For more information, including rationale, tips and tricks, see [Modern Conda Environments with Pixi](./modern-conda-environments-with-pixi.md).
 

diff --git a/openfold3/tests/conftest.py b/openfold3/tests/conftest.py
@@ -138,6 +138,11 @@ def _check_snapshot_env(snapshot_dir: Path) -> None:
         )
 
 
+def _snapshot_platform() -> str:
+    """Return 'rocm' when running on an AMD GPU, 'nvidia' otherwise."""
+    return "rocm" if torch.version.hip is not None else "nvidia"
+
+
 def _write_snapshot_env(snapshot_dir: Path) -> None:
     """Write current environment metadata alongside snapshots."""
     SnapshotEnv.current().to_json(snapshot_dir / _SNAPSHOT_ENV_FILE)
@@ -150,8 +155,9 @@ def pytest_sessionfinish(session, exitstatus):
     snapshots_root = Path(__file__).parent / "test_data" / "snapshots"
     if snapshots_root.exists():
         for subdir in snapshots_root.iterdir():
-            if subdir.is_dir() and any(subdir.glob("*.npz")):
-                _write_snapshot_env(subdir)
+            platform_dir = subdir / _snapshot_platform()
+            if platform_dir.is_dir() and any(platform_dir.glob("*.npz")):
+                _write_snapshot_env(platform_dir)
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -234,9 +240,13 @@ def biotite_ccd_wrapper():
 
 @pytest.fixture(scope="module")
 def original_datadir(request: pytest.FixtureRequest) -> Path:
-    """Redirect pytest-regressions snapshot storage to test_data/snapshots/."""
+    """Redirect pytest-regressions snapshot storage to test_data/snapshots/<platform>/."""
     datadir = (
-        Path(__file__).parent / "test_data" / "snapshots" / Path(request.path).stem
+        Path(__file__).parent
+        / "test_data"
+        / "snapshots"
+        / Path(request.path).stem
+        / _snapshot_platform()
     )
     _check_snapshot_env(datadir)
     return datadir

diff --git a/...t_triangular_attention/_snapshot_env.json → ...gular_attention/nvidia/_snapshot_env.json b/...t_triangular_attention/_snapshot_env.json → ...gular_attention/nvidia/_snapshot_env.json
diff --git a/...gular_attention/test_shape_cpu_False_.npz → ...ttention/nvidia/test_shape_cpu_False_.npz b/...gular_attention/test_shape_cpu_False_.npz → ...ttention/nvidia/test_shape_cpu_False_.npz
diff --git a/...ngular_attention/test_shape_cpu_True_.npz → ...attention/nvidia/test_shape_cpu_True_.npz b/...ngular_attention/test_shape_cpu_True_.npz → ...attention/nvidia/test_shape_cpu_True_.npz
diff --git a/...ular_attention/test_shape_cuda_False_.npz → ...tention/nvidia/test_shape_cuda_False_.npz b/...ular_attention/test_shape_cuda_False_.npz → ...tention/nvidia/test_shape_cuda_False_.npz
diff --git a/...gular_attention/test_shape_cuda_True_.npz → ...ttention/nvidia/test_shape_cuda_True_.npz b/...gular_attention/test_shape_cuda_True_.npz → ...ttention/nvidia/test_shape_cuda_True_.npz
diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/_snapshot_env.json b/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/_snapshot_env.json
@@ -0,0 +1,7 @@
+{
+  "torch_version": "2.11.0+rocm7.2",
+  "python_version": "3.13.13",
+  "cuda_version": null,
+  "cudnn_version": "3005001",
+  "gpu_name": "AMD Instinct MI300X"
+}
diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cpu_False_.npz b/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cpu_False_.npz
diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cpu_True_.npz b/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cpu_True_.npz
diff --git a/...fold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cuda_False_.npz b/...fold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cuda_False_.npz
diff --git a/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cuda_True_.npz b/openfold3/tests/test_data/snapshots/test_triangular_attention/rocm/test_shape_cuda_True_.npz
diff --git a/..._multiplicative_update/_snapshot_env.json → ...licative_update/nvidia/_snapshot_env.json b/..._multiplicative_update/_snapshot_env.json → ...licative_update/nvidia/_snapshot_env.json
diff --git a/...multiplicative_update/test_shape_cpu_.npz → ...icative_update/nvidia/test_shape_cpu_.npz b/...multiplicative_update/test_shape_cpu_.npz → ...icative_update/nvidia/test_shape_cpu_.npz
diff --git a/...ultiplicative_update/test_shape_cuda_.npz → ...cative_update/nvidia/test_shape_cuda_.npz b/...ultiplicative_update/test_shape_cuda_.npz → ...cative_update/nvidia/test_shape_cuda_.npz
diff --git a/...3/tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/_snapshot_env.json b/...3/tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/_snapshot_env.json
@@ -0,0 +1,7 @@
+{
+  "torch_version": "2.11.0+rocm7.2",
+  "python_version": "3.13.13",
+  "cuda_version": null,
+  "cudnn_version": "3005001",
+  "gpu_name": "AMD Instinct MI300X"
+}
diff --git a/.../tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/test_shape_cpu_.npz b/.../tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/test_shape_cpu_.npz
diff --git a/...tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/test_shape_cuda_.npz b/...tests/test_data/snapshots/test_triangular_multiplicative_update/rocm/test_shape_cuda_.npz
diff --git a/pixi.lock b/pixi.lock
diff --git a/pixi.toml b/pixi.toml
@@ -133,6 +133,14 @@ openfold3-cuda13-pypi = { no-default-feature = true, features = [
     "openfold3-editable-full",
     "tests"
 ] }
+openfold3-rocm7 = { no-default-feature = true, features = [
+    "python",
+    "openfold3-cpu",
+    "pytorch-pypi-rocm7",
+    "not-in-pypi",
+    "openfold3-editable",
+    "tests", 
+] }
 
 # Maybe we should have solve-groups, but how to conciliate between e.g. cu12 and cu13?
 # N.B. at the moment all these are editable installs
@@ -250,6 +258,18 @@ torch = { version = "*", index = "https://download.pytorch.org/whl/cu129" }
 pytorch-lightning = "*"
 torch = { version = "*", index = "https://download.pytorch.org/whl/cu130" }
 
+[feature.pytorch-pypi-rocm7]
+
+  platforms = ["linux-64"]
+
+  [feature.pytorch-pypi-rocm7.dependencies]
+  # torch 2.11.0+rocm7.2 requires setuptools<82
+  setuptools = "<82"
+
+  [feature.pytorch-pypi-rocm7.pypi-dependencies]
+  torch = { version = "*", index = "https://download.pytorch.org/whl/rocm7.2" }
+  triton-rocm = { version = "*", index = "https://download.pytorch.org/whl/rocm7.2" }
+
 [feature.cuda12]
 
   platforms = ["linux-64", "linux-aarch64"]