python3Packages.quack-kernels: init at 0.5.0

Co-authored-by: Sizhe Zhao <prc.zhao@outlook.com>

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
+    ];
+  };
+})

pkgs/top-level/python-packages.nix

@@ -16680,6 +16680,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 { };