feat: improve tests and bump build

build.toml

@@ -1,16 +1,18 @@
 [general]
-version = "0.0.1"
+name = "mel_spectrogram"
 
 [torch]
-name = "mel_spectrogram"
 src = [
-  "ext-torch/registration.h",
-  "ext-torch/torch_binding.cpp",
-  "ext-torch/torch_binding.h",
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h"
 ]
-pyroot = "ext-torch"
 
 [kernel.mel_spectrogram]
-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
-src = ["mel_spectrogram/mel_spectrogram.cu"]
-depends = ["torch"]
+cuda-capabilities = [ 
+   # "7.0", "7.2", "7.5", 
+   "8.0", "8.6", "8.7", "8.9", "9.0" 
+]
+src = [
+  "mel_spectrogram/mel_spectrogram.cu"
+]
+depends = [ "torch" ]

build/torch24-cxx11-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch24-cxx11-cu121-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch24-cxx11-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch24-cxx98-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch24-cxx98-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d53f5e8ce2a7dae51f5481b23a378ec0b71b681705f77b43626fc66980e26751
-size 56640

build/torch24-cxx98-cu121-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch24-cxx98-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:891c8c25a5d4ccdd25458f940d83e21665fb8d1936a6990fb169420dc375dea7
-size 52216

build/torch24-cxx98-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +0,0 @@
-import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1

build/torch25-cxx11-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:93bc59bb2ba5e7a1bdccf6cf76bf9a902a9fe5db14de1114f343ac3a7e1efc78
-size 56832

build/{torch24-cxx11-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so → torch25-cxx11-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2d23c602fb3ff07cd05745af022156bd7fea563c9e7e85cda47ce5b59f49910e
-size 56832
+oid sha256:c9fd70ede706f4c301b0a956581e1eca47fd59ffe151b02bae892acaedff1119
+size 2019488

build/torch25-cxx11-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch25-cxx11-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:93cbaec30c56f75ef21a299cf551fcbc4bc66cd1a9b32560c1fddaf403f093d7
-size 52328

build/{torch24-cxx11-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so → torch25-cxx11-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:86654dcf62b8e5a226eab8191e156368cbafe973eda7e719efdf628002cda52d
-size 52120
+oid sha256:8be112fecfd9abc4fd2916d6352e5b5ce1d327d16966faddf7739c5f7a812368
+size 2042320

build/torch25-cxx11-cu121-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch25-cxx11-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:bdb6dfbb056ca9f0ab604f4d18558dfec95fa8df6968b40eb3ac40bfda413912
-size 52120

build/{torch24-cxx11-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so → torch25-cxx11-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c6d28cf590e366a96f1a9b68fdc73a52acd909a6c2c828e0e7724ce0ed4bfbc4
-size 52328
+oid sha256:88b0a4333c16b4c8e0e307f2cdfa4461e3a2f0a7a00007f81193c5f4bfbdc658
+size 2056320

build/torch25-cxx11-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch25-cxx98-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e9c16679793bf509d85bbd2f58d63ef7235b07c617fe44189781a2622cca9a7b
-size 57128

build/{torch24-cxx98-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so → torch25-cxx98-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c7b04d9d71064fc2219d5bc522caec86b07f7b3b73a72894c6221315b7833914
-size 57128
+oid sha256:3675a983b7da77eeaa944185ad5b579687866f38d254c3c0be9207473c4f1733
+size 2012128

build/torch25-cxx98-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch25-cxx98-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ee2b160940af6bb144e78c37a91b365ad6744a266b68a085a772a6d94dec2ae8
-size 56640

build/torch25-cxx98-cu121-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b67c2dbe9bdb3a6c8f8eb1bcc7ece574467db1697ee95e273c9addc6a3bc601b
+size 2039192

build/torch25-cxx98-cu121-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch25-cxx98-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_0_0_1.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:e4cfd1fc6fc2210c286dd909951b8554262646aba8382f8a9bbb1e0db4819f44
-size 52216

build/torch25-cxx98-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:289311e0fc29f6dfa7561b4d5bd83d7e43330e6d2fc8e3d7242b516a63eb94eb
+size 2053192

build/torch25-cxx98-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -1,3 +1,9 @@
 import torch
-from . import _mel_spectrogram_0_0_1
-ops = torch.ops._mel_spectrogram_0_0_1
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d0ed57bdcc97d6cbdf4b9002e1f6c9939aaa911ef5a69ffaf89682ce473feb7
+size 2019712

build/torch26-cxx11-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:804671cc1e390382553d29afbe8daf05a65e4cf8b6785d1585e5d791b91d8c10
+size 2056536

build/torch26-cxx11-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ff08507c8dac97d4f9dbb817cba14a35a48df9f63a5be1cf4b3c02fd4d6cc44
+size 2061080

build/torch26-cxx11-cu126-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0dfe82be4faeafacd152915330f041edec27b426604feca191075b0889850cd8
+size 2012400

build/torch26-cxx98-cu118-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a455e5d5c93bd47127c9061e9b1084625e0405bee8bd6958abb559998a62bc28
+size 2053408

build/torch26-cxx98-cu124-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/mel_spectrogram/_mel_spectrogram_dff20e5_dirty.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5a10c07449a64d5d58d1126778d9a620bf74abf227dbb37a1cde9d06236f7cd
+size 2057952

build/torch26-cxx98-cu126-x86_64-linux/mel_spectrogram/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _mel_spectrogram_dff20e5_dirty
+ops = torch.ops._mel_spectrogram_dff20e5_dirty
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_mel_spectrogram_dff20e5_dirty::{op_name}"

ext-torch/torch_binding.h

@@ -1,10 +0,0 @@
-#pragma once
-
-#include <torch/torch.h>
-
-void mel_spectrogram(
-    torch::Tensor &out,
-    torch::Tensor const &samples,
-    torch::Tensor const &filters,
-    int64_t fft_size,
-    int64_t fft_step);

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1733328505,
+        "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rocm-nix": "rocm-nix"
+      },
+      "locked": {
+        "lastModified": 1743422493,
+        "narHash": "sha256-3RBEo/vw1p1CMGHsZEyP9Z9gRmJN+USoDMCSVlxHGXg=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "2170048e94663eb8012e62a84b0ba9c1a637e2b3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1740557110,
+        "narHash": "sha256-D2waFyJkaepTchTrGVAIfCd/YP+37bgXWg9cXwuxuT0=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "b89a821293c3872992137114d0db9a791243a41b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "rocm-nix": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1743085847,
+        "narHash": "sha256-uWG29p+nhZmGRV1LffWwRGjwtPIXeu1F0YTQbXgB+GU=",
+        "owner": "huggingface",
+        "repo": "rocm-nix",
+        "rev": "245cdc9bfb4bfafa818711c5f5e0b889afe1ba39",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "rocm-nix",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -2,7 +2,7 @@
   description = "Flake for mel_spectogram";
 
   inputs = {
-    kernel-builder.url = "git+ssh://git@github.com/huggingface/kernel-builder";
+    kernel-builder.url = "github:huggingface/kernel-builder";
   };
 
   outputs =

mel_spectrogram/mel_spectrogram.cu

@@ -100,9 +100,9 @@ __global__ void normalize_mel_spec_kernel(
 }
 
 void mel_spectrogram(
-    torch::Tensor &out,        // [n_mel, n_frames]
-    torch::Tensor const &samples,    // [n_samples]
-    torch::Tensor const &filters,    // [n_mel, n_fft]
+    const torch::Tensor &out,        // [n_mel, n_frames]
+    const torch::Tensor &samples,    // [n_samples]
+    const torch::Tensor &filters,    // [n_mel, n_fft]
     const int64_t fft_size,
     const int64_t fft_step) {