feat: Add torch stub

README.md

@@ -10,3 +10,10 @@ Execute
 $ nix build
 $ result/bin/reproducible
 ```
+
+## Troubleshooting
+
+Resolve collision:
+``` sh
+sudo nix-store --verify --check-contents --repair
+```

flake.lock

@@ -59,11 +59,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1710631334,
+        "lastModified": 1711163522,
-        "narHash": "sha256-rL5LSYd85kplL5othxK5lmAtjyMOBg390sGBTb3LRMM=",
+        "narHash": "sha256-YN/Ciidm+A0fmJPWlHBGvVkcarYWSC+s3NTPk/P+q3c=",
         "owner": "nixos",
         "repo": "nixpkgs",
-        "rev": "c75037bbf9093a2acb617804ee46320d6d1fea5a",
+        "rev": "44d0940ea560dee511026a53f0e2e2cde489b4d4",
         "type": "github"
       },
       "original": {

flake.nix

@@ -23,39 +23,38 @@
       "x86_64-linux"
       "x86_64-darwin"
     ];
-    perSystem = { system, pkgs, ... }: let
+    perSystem = { system, ... }: let
-      inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; })
+      pkgs = import nixpkgs {
-        defaultPoetryOverrides mkPoetryApplication mkPoetryEnv;
+        inherit system;
-      args = {
+        config = {
-        python = pkgs.python311;
+          allowUnfree = true;
-
+          cudaSupport = true;
-        # Required since pycapnp can't be built for some reason - Maybe this is a
-        # missing dependency issue.
-        preferWheels = true;
-        overrides = defaultPoetryOverrides;
         };
-      application = mkPoetryApplication (args // {
+      };
-        projectDir = ./.;
+      poetryLib = poetry2nix.lib.mkPoetry2Nix { inherit pkgs; };
-      });
+      common = {
-      env = mkPoetryEnv (args // {
+        python = pkgs.python311;
         projectDir = ./.;
+        preferWheels = true;
+        overrides = poetryLib.defaultPoetryOverrides;
+      };
+      application = poetryLib.mkPoetryApplication common;
+      env = poetryLib.mkPoetryEnv (common // {
         editablePackageSources = let project_dir = builtins.getEnv "PROJECT_DIR"; in
         {
           reproducible = if project_dir == "" then ./. else "${project_dir}";
         };
       });
-      dockerImage = pkgs.dockerTools.buildImage {
-        name = "reproducible";
-        config = { Cmd = [ "python" "main.py" ]; };
-      };
     in rec {
       packages = {
         default = application;
-        inherit dockerImage;
       };
       devShells.default = pkgs.mkShell {
-        packages = [ env pkgs.poetry ];
+        packages = [ env ];
       };
+      #devShells.default = env.overrideAttrs (oldAttrs: {
+      #  buildInputs = [ args.python ];
+      #});
     };
   };
 }

pyproject.toml

@@ -7,13 +7,42 @@ readme = "README.md"
 include = []
 packages = [{ include = "reproducible" }]
 
-[tool.poetry.dependencies]
-python = "^3.11"
-numpy = "^1.26"
-
 [tool.poetry.scripts]
-reproducible = "reproducible.entry:main"
+reproducible = "reproducible.util:main"
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.dependencies]
+python = "^3.11"
+numpy = "^1.26"
+#torch = { url = "https://download.pytorch.org/whl/cu118/torch-2.2.1+cu118-cp311-cp311-linux_x86_64.whl" }
+torch = { version = "^2.2.0", source = "pytorch-cu118" }
+torch-geometric = "^2.5.2"
+#pyg-lib = { git = "https://github.com/pyg-team/pyg-lib.git", tag = "0.4.0" }
+#torch-cluster = { version = "1.6.3" }
+#torch-scatter = { version = "2.1.2" }
+#torch-sparse = { version = "0.6.18" }
+#torch-spline-conv = { version = "1.2.2" }
+
+#pyg-lib = { version = "0.4.0", source = "pyg-torch220-cu118" }
+#torch-cluster = { version = "1.6.3", source = "pyg-torch220-cu118" }
+#torch-scatter = { version = "2.1.2", source = "pyg-torch220-cu118" }
+#torch-sparse = { version = "0.6.18", source = "pyg-torch220-cu118" }
+#torch-spline-conv = { version = "1.2.2", source = "pyg-torch220-cu118" }
+
+[[tool.poetry.source]]
+name = "pytorch-cu118"
+url = "https://download.pytorch.org/whl/cu118"
+priority = "explicit"
+
+#[[tool.poetry.source]]
+#name = "pyg-torch220-cu118"
+#url = "git+https://github.com/pyg-team/pyg-lib.git"
+#url = "https://data.pyg.org/whl/torch-2.2.0+cu118.html"
+#priority = "explicit"
+
+#[[tool.poetry.source]]
+#name = "PyPI"
+#priority = "primary"

reproducible/__main__.py

@@ -1,3 +1,3 @@
-import reproducible.entry
+import reproducible.util
 
-reproducible.entry.main()
+reproducible.util.main()

reproducible/entry.py

@@ -1,2 +0,0 @@
-def main():
-    print("Hi")

reproducible/util.py

@@ -0,0 +1,8 @@
+import torch
+import torch_geometric
+
+def main():
+    print(f"Torch version: {torch.__version__}")
+    print(f"CUDA version: {torch.version.cuda}")
+    print(f"Cuda is available? {torch.cuda.is_available()}")
+    print(f"PyG version: {torch_geometric.__version__}")