From 2cc58226b0205fae68ff9266aa6edef09df06f6d Mon Sep 17 00:00:00 2001
From: Someone Serge <sergei.kozlukov@aalto.fi>
Date: Sat, 20 Jul 2024 12:05:52 +0000
Subject: [PATCH] python3Packages.torch.tests.*compile*: init

---
 .../torch/mk-torch-compile-check.nix          | 38 +++++++++++++++++++
 .../python-modules/torch/tests.nix            | 16 +++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 pkgs/development/python-modules/torch/mk-torch-compile-check.nix

diff --git a/pkgs/development/python-modules/torch/mk-torch-compile-check.nix b/pkgs/development/python-modules/torch/mk-torch-compile-check.nix
new file mode 100644
index 000000000000..268ed5297da9
--- /dev/null
+++ b/pkgs/development/python-modules/torch/mk-torch-compile-check.nix
@@ -0,0 +1,38 @@
+{
+  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
+
+
+    @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})))
+  ''
diff --git a/pkgs/development/python-modules/torch/tests.nix b/pkgs/development/python-modules/torch/tests.nix
index 92ddccccbaeb..e3f2ca44ba5a 100644
--- a/pkgs/development/python-modules/torch/tests.nix
+++ b/pkgs/development/python-modules/torch/tests.nix
@@ -1,6 +1,6 @@
 { callPackage }:
 
-{
+rec {
   # To perform the runtime check use either
   # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or
   # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox)
@@ -14,4 +14,18 @@
     versionAttr = "hip";
     libraries = ps: [ ps.torchWithRocm ];
   };
+
+  compileCpu = tester-compileCpu.gpuCheck;
+  tester-compileCpu = callPackage ./mk-torch-compile-check.nix {
+    feature = null;
+    libraries = ps: [ ps.torch ];
+  };
+  tester-compileCuda = callPackage ./mk-torch-compile-check.nix {
+    feature = "cuda";
+    libraries = ps: [ ps.torchWithCuda ];
+  };
+  tester-compileRocm = callPackage ./mk-torch-compile-check.nix {
+    feature = "rocm";
+    libraries = ps: [ ps.torchWithRocm ];
+  };
 }