diff --git a/pkgs/development/python-modules/quack-kernels/default.nix b/pkgs/development/python-modules/quack-kernels/default.nix new file mode 100644 index 000000000000..b817693d6d55 --- /dev/null +++ b/pkgs/development/python-modules/quack-kernels/default.nix @@ -0,0 +1,75 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + + # dependencies + apache-tvm-ffi, + einops, + nvidia-cutlass-dsl, + torch, + torch-c-dlpack-ext, + + # optional-dependencies + # nvidia-matmul-heuristics, + jax, + # jax-tvm-ffi, + pandas, +}: +buildPythonPackage (finalAttrs: { + pname = "quack-kernels"; + version = "0.5.0"; + pyproject = true; + __structuredAttrs = true; + + src = fetchFromGitHub { + owner = "Dao-AILab"; + repo = "quack"; + tag = "v${finalAttrs.version}"; + hash = "sha256-Y56jJUTn/HopOe0yNxxxwMf+abXSdzTa8+YoiHd/rFE="; + }; + + build-system = [ + setuptools + ]; + + dependencies = [ + apache-tvm-ffi + einops + nvidia-cutlass-dsl + torch + torch-c-dlpack-ext + ]; + + optional-dependencies = { + heuristics = [ + # nvidia-matmul-heuristics + ]; + jax = [ + jax + # jax-tvm-ffi + ]; + bench = [ + pandas + ]; + }; + + pythonImportsCheck = [ "quack" ]; + + # Fatal Python error: Aborted + doCheck = false; + + meta = { + description = "Quirky Assortment of CuTe Kernels"; + homepage = "https://github.com/Dao-AILab/quack"; + changelog = "https://github.com/Dao-AILab/quack/releases/tag/${finalAttrs.src.tag}"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ + GaetanLepage + prince213 + ]; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index b18a21f93d76..fec3053e53d8 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -16682,6 +16682,8 @@ self: super: with self; { qtpy = callPackage ../development/python-modules/qtpy { }; + quack-kernels = callPackage ../development/python-modules/quack-kernels { }; + quadprog = callPackage ../development/python-modules/quadprog { }; qualysclient = callPackage ../development/python-modules/qualysclient { };