tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 cudaPackages,
3 feature ? null,
4 lib,
5 libraries,
6 name ? if feature == null then "torch-compile-cpu" else "torch-compile-${feature}",
7 stdenv,
8 writableTmpDirAsHomeHook,
9}:
10let
11 deviceStr = if feature == null then "" else '', device="cuda"'';
12in
13cudaPackages.writeGpuTestPython
14 {
15 inherit name feature libraries;
16
17 # This could be accomplished with propagatedBuildInputs instead of introducing
18 gpuCheckArgs.nativeBuildInputs = [
19 # torch._inductor.exc.InductorError: PermissionError: [Errno 13] Permission denied: '/homeless-shelter'
20 writableTmpDirAsHomeHook
21 ];
22 makeWrapperArgs = [
23 "--suffix"
24 "PATH"
25 ":"
26 "${lib.getBin stdenv.cc}/bin"
27 ];
28 }
29 ''
30 import torch
31
32
33 @torch.compile
34 def opt_foo2(x, y):
35 a = torch.sin(x)
36 b = torch.cos(y)
37 return a + b
38
39
40 print(
41 opt_foo2(
42 torch.randn(10, 10${deviceStr}),
43 torch.randn(10, 10${deviceStr})))
44 ''