push sheeet

pkgs/development/libraries/science/math/libtorch/test/CMakeLists.txt

@@ -0,0 +1,5 @@
+cmake_minimum_required(VERSION 3.0)
+find_package(Torch REQUIRED)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+add_executable(test test.cpp)
+target_link_libraries(test "${TORCH_LIBRARIES}")

pkgs/development/libraries/science/math/libtorch/test/default.nix

@@ -0,0 +1,64 @@
+{
+  lib,
+  stdenv,
+  cmake,
+  libtorch-bin,
+  linkFarm,
+  symlinkJoin,
+
+  cudaSupport,
+  cudaPackages ? { },
+}:
+let
+  inherit (cudaPackages) cudatoolkit cudnn;
+
+  cudatoolkit_joined = symlinkJoin {
+    name = "${cudatoolkit.name}-unsplit";
+    paths = [
+      cudatoolkit.out
+      cudatoolkit.lib
+    ];
+  };
+
+  # We do not have access to /run/opengl-driver/lib in the sandbox,
+  # so use a stub instead.
+  cudaStub = linkFarm "cuda-stub" [
+    {
+      name = "libcuda.so.1";
+      path = "${cudatoolkit}/lib/stubs/libcuda.so";
+    }
+  ];
+
+in
+stdenv.mkDerivation {
+  pname = "libtorch-test";
+  version = libtorch-bin.version;
+
+  src = lib.fileset.toSource {
+    root = ./.;
+    fileset = lib.fileset.unions [
+      ./CMakeLists.txt
+      ./test.cpp
+    ];
+  };
+
+  nativeBuildInputs = [ cmake ];
+
+  buildInputs = [ libtorch-bin ] ++ lib.optionals cudaSupport [ cudnn ];
+
+  cmakeFlags = lib.optionals cudaSupport [ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
+
+  doCheck = true;
+
+  installPhase = ''
+    touch $out
+  '';
+
+  checkPhase =
+    lib.optionalString cudaSupport ''
+      LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
+    ''
+    + ''
+      ./test
+    '';
+}

pkgs/development/libraries/science/math/libtorch/test/test.cpp

@@ -0,0 +1,20 @@
+#undef NDEBUG
+#include <cassert>
+
+#include <iostream>
+
+#include <torch/torch.h>
+
+int main() {
+  torch::Tensor tensor = torch::eye(3);
+
+  float checkData[] = {
+    1, 0, 0,
+    0, 1, 0,
+    0, 0, 1
+  };
+
+  torch::Tensor check = torch::from_blob(checkData, {3, 3});
+
+  assert(tensor.allclose(check));
+}