python3Packages.torch.gpuChecks: add rocm

pkgs/development/python-modules/torch/default.nix

@@ -648,7 +648,7 @@ buildPythonPackage rec {
     blasProvider = blas.provider;
     # To help debug when a package is broken due to CUDA support
     inherit brokenConditions;
-    gpuChecks.cudaAvailable = callPackage ./test-cuda.nix { torch = torchWithCuda; };
+    gpuChecks = callPackage ./gpu-checks.nix { };
   };
 
   meta = {

pkgs/development/python-modules/torch/gpu-checks.nix

@@ -0,0 +1,50 @@
+{
+  lib,
+  callPackage,
+  torchWithCuda,
+  torchWithRocm,
+}:
+
+let
+  accelAvailable =
+    {
+      feature,
+      versionAttr,
+      torch,
+      runCommandNoCC,
+      writers,
+    }:
+    let
+      name = "${torch.name}-${feature}-check";
+      unwrapped = writers.writePython3Bin "${name}-unwrapped" {libraries = [torch];} ''
+        import torch
+        message = f"{torch.cuda.is_available()=} and {torch.version.${versionAttr}=}"
+        assert torch.cuda.is_available() and torch.version.${versionAttr}, message
+        print(message)
+      '';
+    in
+    runCommandNoCC name
+      {
+        nativeBuildInputs = [unwrapped];
+        requiredSystemFeatures = [feature];
+        passthru = {
+          inherit unwrapped;
+        };
+      }
+      ''
+        ${name}-unwrapped
+        touch $out
+      '';
+in
+{
+  cudaAvailable = callPackage accelAvailable {
+    feature = "cuda";
+    versionAttr = "cuda";
+    torch = torchWithCuda;
+  };
+  rocmAvailable = callPackage accelAvailable {
+    feature = "rocm";
+    versionAttr = "hip";
+    torch = torchWithRocm;
+  };
+}

pkgs/development/python-modules/torch/test-cuda.nix

@@ -1,21 +0,0 @@
-{ runCommandNoCC
-, python
-, torch
-}:
-
-runCommandNoCC "${torch.name}-gpu-test"
-{
-  nativeBuildInputs = [
-    (python.withPackages (_: [ torch ]))
-  ];
-  requiredSystemFeatures = [
-    "cuda"
-  ];
-} ''
-  python3 << EOF
-  import torch
-  assert torch.cuda.is_available(), f"{torch.cuda.is_available()=}"
-  EOF
-
-  touch $out
-''