python3Packages.torch.tests: tests for darwin

Run nixfmt for mk-torch-compile-check.nix
On darwin, both CPU and MPS are tested now.

Signed-off-by: Mika Tammi <[email protected]>

pkgs/development/python-modules/torch/mk-torch-compile-check.nix

@@ -1,38 +1,34 @@
-{
-  cudaPackages,
-  feature ? null,
-  lib,
-  libraries,
-  name ? if feature == null then "torch-compile-cpu" else "torch-compile-${feature}",
-  pythonPackages,
-  stdenv,
-}:
-let
-  deviceStr = if feature == null then "" else '', device="cuda"'';
-in
-(cudaPackages.writeGpuTestPython.override { python3Packages = pythonPackages; })
-  {
-    inherit name feature libraries;
-    makeWrapperArgs = [
-      "--suffix"
-      "PATH"
-      ":"
-      "${lib.getBin stdenv.cc}/bin"
-    ];
-  }
-  ''
-    import torch
+{ cudaPackages, feature ? null, lib, libraries, name ?
+  if feature == null then "torch-compile-cpu" else "torch-compile-${feature}"
+, pythonPackages, stdenv, }:
+let deviceStr = if feature == null then "" else '', device="${feature}"'';
+in (cudaPackages.writeGpuTestPython.override {
+  python3Packages = pythonPackages;
+}) {
+  inherit name feature libraries;
+  makeWrapperArgs = [ "--suffix" "PATH" ":" "${lib.getBin stdenv.cc}/bin" ];
+} (''
+  import torch
 
 
-    @torch.compile
-    def opt_foo2(x, y):
-        a = torch.sin(x)
-        b = torch.cos(y)
-        return a + b
+'' + lib.optionalString (!stdenv.hostPlatform.isDarwin) ''
+  # torch.compile requires OpenMP which is not available on Darwin
+  @torch.compile
+'' + ''
+  def opt_foo2(x, y):
+      a = torch.sin(x)
+      b = torch.cos(y)
+      return a + b
 
 
-    print(
-      opt_foo2(
-        torch.randn(10, 10${deviceStr}),
-        torch.randn(10, 10${deviceStr})))
-  ''
+  print(
+    opt_foo2(
+      torch.randn(10, 10${deviceStr}),
+      torch.randn(10, 10${deviceStr})))
+'' + lib.optionalString stdenv.hostPlatform.isDarwin ''
+  # MPS is built by default on Darwin, so test it.
+  print(
+    opt_foo2(
+      torch.randn(10, 10, device="mps"),
+      torch.randn(10, 10, device="mps")))
+'')