diff --git a/pkgs/development/python-modules/torch-c-dlpack-ext/default.nix b/pkgs/development/python-modules/torch-c-dlpack-ext/default.nix new file mode 100644 index 000000000000..07d254591587 --- /dev/null +++ b/pkgs/development/python-modules/torch-c-dlpack-ext/default.nix @@ -0,0 +1,40 @@ +{ + lib, + buildPythonPackage, + apache-tvm-ffi, + + # build-system + setuptools, + + # dependencies + torch, +}: + +buildPythonPackage (finalAttrs: { + pname = "torch-c-dlpack-ext"; + inherit (apache-tvm-ffi) version src; + pyproject = true; + + sourceRoot = "${finalAttrs.src.name}/addons/torch_c_dlpack_ext"; + + build-system = [ + apache-tvm-ffi + setuptools + ]; + + dependencies = [ + torch + ]; + + pythonImportsCheck = [ "torch_c_dlpack_ext" ]; + + # No tests + doCheck = false; + + meta = { + description = "Ahead-Of-Time (AOT) compiled module to support faster DLPack conversion in DLPack"; + homepage = "https://github.com/apache/tvm-ffi/tree/main/addons/torch_c_dlpack_ext"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ GaetanLepage ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 94a3b02b205b..72c71514d080 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -19453,6 +19453,8 @@ self: super: with self; { torch-bin = callPackage ../development/python-modules/torch/bin { triton = self.triton-bin; }; + torch-c-dlpack-ext = callPackage ../development/python-modules/torch-c-dlpack-ext { }; + torch-geometric = callPackage ../development/python-modules/torch-geometric { }; # Required to test triton