auvux-dsp 0.1.0.dev0__tar.gz → 0.1.0.dev2__tar.gz

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/.gitignore

@@ -9,3 +9,8 @@ __pycache__/
 .mypy_cache/
 .ruff_cache/
 .DS_Store
+.env
+.env.*
+*.env
+.pypirc
+*.pem

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/CMakeLists.txt

@@ -49,6 +49,11 @@ pybind11_add_module(_native
   src/ops/cqt/cqt_filterbank.cpp
   src/ops/cqt/cqt_cpu.cpp
   src/ops/chroma/chroma_cpu.cpp
+  src/ops/moments/moments_cpu.cpp
+  src/ops/rms/rms_cpu.cpp
+  src/ops/flux/flux_cpu.cpp
+  src/ops/onset/onset_cpu.cpp
+  src/ops/tempogram/tempogram_cpu.cpp
   src/gpu/gpu_stub.cpp
   src/gpu/cqt_plan.cpp
   src/third_party/pffft.c
@@ -57,6 +62,7 @@ pybind11_add_module(_native
   src/bindings/bind_stft.cpp
   src/bindings/bind_mel.cpp
   src/bindings/bind_cqt.cpp
+  src/bindings/bind_features.cpp
   src/bindings/bind_util.cpp)
 target_include_directories(_native PRIVATE src src/third_party)
 target_link_libraries(_native PRIVATE Threads::Threads)
@@ -90,7 +96,12 @@ if(AUVUX_GPU STREQUAL "metal")
       ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/common.metal
       ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/stft.metal
       ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/mel.metal
-      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/cqt.metal)
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/cqt.metal
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/rms.metal
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/moments.metal
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/flux.metal
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/onset.metal
+      ${CMAKE_CURRENT_SOURCE_DIR}/src/gpu/metal/kernels/tempogram.metal)
   set(auvux_metal_gen ${CMAKE_CURRENT_BINARY_DIR}/generated/auvux_metal_kernels.h)
   string(JOIN "," auvux_metal_kernels_arg ${auvux_metal_kernels})
   add_custom_command(
@@ -104,7 +115,12 @@ if(AUVUX_GPU STREQUAL "metal")
       src/gpu/metal/metal_common.mm
       src/gpu/metal/metal_stft.mm
       src/gpu/metal/metal_mel.mm
-      src/gpu/metal/metal_cqt.mm)
+      src/gpu/metal/metal_cqt.mm
+      src/gpu/metal/metal_rms.mm
+      src/gpu/metal/metal_moments.mm
+      src/gpu/metal/metal_flux.mm
+      src/gpu/metal/metal_onset.mm
+      src/gpu/metal/metal_tempogram.mm)
   target_sources(_native PRIVATE ${auvux_metal_srcs} ${auvux_metal_gen})
   set_source_files_properties(${auvux_metal_srcs} PROPERTIES COMPILE_OPTIONS "-fobjc-arc")
   target_include_directories(_native PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/generated)
@@ -125,7 +141,8 @@ elseif(AUVUX_GPU STREQUAL "cuda")
     src/gpu/cuda/cuda_common.cu
     src/gpu/cuda/cuda_stft.cu
     src/gpu/cuda/cuda_mel.cu
-    src/gpu/cuda/cuda_cqt.cu)
+    src/gpu/cuda/cuda_cqt.cu
+    src/gpu/cuda/cuda_features.cu)
   target_compile_definitions(_native PRIVATE AUVUX_CUDA=1)
   if(NOT auvux_user_archs)
     # Distributable fatbin: SASS per supported generation plus PTX (the plain

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: auvux-dsp
-Version: 0.1.0.dev0
+Version: 0.1.0.dev2
 Summary: Fast differentiable audio transforms (STFT, mel, MFCC, CQT, chroma) on CPU and GPU
 Keywords: audio,dsp,stft,mel,mfcc,cqt,chroma,spectrogram,gpu
 Author-Email: Peter Kiers <pkiers.1983@gmail.com>
@@ -17,15 +17,18 @@ Requires-Dist: pytest; extra == "test"
 Requires-Dist: librosa; extra == "test"
 Requires-Dist: soxr; extra == "test"
 Requires-Dist: torch; extra == "test"
+Provides-Extra: wheeltest
+Requires-Dist: pytest; extra == "wheeltest"
 Description-Content-Type: text/markdown
 
 # auvux-dsp
 
 Fast differentiable audio transforms (STFT, iSTFT, mel, MFCC, CQT/VQT, chroma)
-with native CPU backends (vDSP on macOS, PFFFT elsewhere) and GPU support
-(Metal, CUDA). Forward *and* backward passes run in native kernels; PyTorch
-autograd plugs in when you pass a torch tensor, and torch is never required
-otherwise.
+plus differentiable scalar/temporal features (spectral moments & flatness, RMS,
+spectral flux, onset strength, tempogram) with native CPU backends (vDSP on
+macOS, PFFFT elsewhere) and GPU support (Metal, CUDA). Forward *and* backward
+passes run in native kernels; PyTorch autograd plugs in when you pass a torch
+tensor, and torch is never required otherwise.
 
 ```python
 import auvux.dsp as dsp
@@ -69,6 +72,9 @@ outputs (power/db/mel/...) are ordinary C-contiguous arrays.
 Status: under construction.
 - CPU (vDSP/PFFFT): STFT, ISTFT, MelSpectrogram, MFCC, CQT, VQT, Chroma —
   forward and native backward, librosa-parity tested, torch autograd built in.
+- Differentiable features: SpectralMoments (centroid, bandwidth, skewness,
+  kurtosis, flatness), RMS, SpectralFlux, OnsetStrength, Tempogram — native
+  forward + analytic adjoint, librosa-parity / torch-grad tested.
 - Metal: all of the above on GPU (n_fft <= 4096), forward + backward,
   parity-tested against the CPU path. torch MPS tensors stay on the GPU end to
   end (DLPack), and backend="auto" routes them there — no flags needed.

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/README.md

@@ -1,10 +1,11 @@
 # auvux-dsp
 
 Fast differentiable audio transforms (STFT, iSTFT, mel, MFCC, CQT/VQT, chroma)
-with native CPU backends (vDSP on macOS, PFFFT elsewhere) and GPU support
-(Metal, CUDA). Forward *and* backward passes run in native kernels; PyTorch
-autograd plugs in when you pass a torch tensor, and torch is never required
-otherwise.
+plus differentiable scalar/temporal features (spectral moments & flatness, RMS,
+spectral flux, onset strength, tempogram) with native CPU backends (vDSP on
+macOS, PFFFT elsewhere) and GPU support (Metal, CUDA). Forward *and* backward
+passes run in native kernels; PyTorch autograd plugs in when you pass a torch
+tensor, and torch is never required otherwise.
 
 ```python
 import auvux.dsp as dsp
@@ -48,6 +49,9 @@ outputs (power/db/mel/...) are ordinary C-contiguous arrays.
 Status: under construction.
 - CPU (vDSP/PFFFT): STFT, ISTFT, MelSpectrogram, MFCC, CQT, VQT, Chroma —
   forward and native backward, librosa-parity tested, torch autograd built in.
+- Differentiable features: SpectralMoments (centroid, bandwidth, skewness,
+  kurtosis, flatness), RMS, SpectralFlux, OnsetStrength, Tempogram — native
+  forward + analytic adjoint, librosa-parity / torch-grad tested.
 - Metal: all of the above on GPU (n_fft <= 4096), forward + backward,
   parity-tested against the CPU path. torch MPS tensors stay on the GPU end to
   end (DLPack), and backend="auto" routes them there — no flags needed.

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/benchmarks/benchmark.py

@@ -1,10 +1,10 @@
-"""Benchmark auvux-dsp against librosa, torchaudio, and nnAudio.
+"""Benchmark auvux-dsp against librosa and torchaudio.
 
 Forward passes and torch training steps (forward + sum().backward()) per
 transform, implementation, and signal length. auvux runs on cpu, host-staged
-gpu, and gpu-resident torch tensors; torchaudio and nnAudio also run on the
-torch device when one is available. Progress prints as it runs; the summary
-tables are markdown, ready for the README.
+gpu, and gpu-resident torch tensors; torchaudio also runs on the torch device
+when one is available. Progress prints as it runs; the summary tables are
+markdown, ready for the README.
 
     python benchmarks/benchmark.py            # everything available
     python benchmarks/benchmark.py --quick    # 1 repeat, short signals
@@ -95,16 +95,33 @@ def build_cases(y, sec, torch, device, gpu, resident):
         ("vqt", vqt_t),
         ("chroma", chroma_t),
     ]
+    # CPU-only differentiable features (no GPU kernels yet). librosa comparisons
+    # are registered below only for the 1:1-matching definitions.
+    feat = [
+        ("spectral_moments", dsp.SpectralMoments(sr=SR, n_fft=2048, hop_length=512)),
+        ("spectral_centroid",
+         dsp.SpectralMoments(sr=SR, features=("centroid",), n_fft=2048, hop_length=512)),
+        ("spectral_bandwidth",
+         dsp.SpectralMoments(sr=SR, features=("bandwidth",), n_fft=2048, hop_length=512)),
+        ("spectral_flatness",
+         dsp.SpectralMoments(sr=SR, features=("flatness",), n_fft=2048, hop_length=512)),
+        ("rms", dsp.RMS(sr=SR, frame_length=2048, hop_length=512)),
+        ("spectral_flux", dsp.SpectralFlux(sr=SR, n_fft=2048, hop_length=512)),
+        ("onset_strength", dsp.OnsetStrength(sr=SR, n_fft=2048, hop_length=512, n_mels=128)),
+        ("tempogram", dsp.Tempogram(sr=SR, n_fft=2048, hop_length=512, n_mels=128)),
+    ]
 
     for name, t in fwd:
         add("forward", name, "auvux cpu", lambda t=t: t(y))
+    for name, t in feat:
+        add("forward", name, "auvux cpu", lambda t=t: t(y))
     if gpu:
-        for name, t in fwd:
+        for name, t in fwd + feat:
             add("forward", name, f"auvux {gpu} staged", lambda t=t: t(y, backend="gpu"))
     if resident:
         sync = syncer(torch, device)
         ytd = torch.from_numpy(y).to(device)
-        for name, t in fwd:
+        for name, t in fwd + feat:
             add(
                 "forward",
                 name,
@@ -168,6 +185,25 @@ def build_cases(y, sec, torch, device, gpu, resident):
             "librosa",
             lambda: librosa.istft(Z, n_fft=2048, hop_length=512, length=n_out),
         )
+        # Features whose definitions match librosa 1:1, so the diff column is a
+        # real correctness check.
+        add("forward", "spectral_centroid", "librosa",
+            lambda: librosa.feature.spectral_centroid(y=y, sr=SR, n_fft=2048, hop_length=512))
+        add("forward", "spectral_bandwidth", "librosa",
+            lambda: librosa.feature.spectral_bandwidth(y=y, sr=SR, n_fft=2048, hop_length=512))
+        add("forward", "spectral_flatness", "librosa",
+            lambda: librosa.feature.spectral_flatness(y=y, n_fft=2048, hop_length=512))
+        add("forward", "rms", "librosa",
+            lambda: librosa.feature.rms(y=y, frame_length=2048, hop_length=512))
+        # Speed-only: our onset/tempogram match librosa's values but not its
+        # `center=True` envelope lag-shift (n_fft//(2*hop) frames), so the diff
+        # is suppressed (see _NO_DIFF) — correctness lives in the test suite.
+        # librosa has no standalone spectral_flux.
+        add("forward", "onset_strength", "librosa",
+            lambda: librosa.onset.onset_strength(
+                y=y, sr=SR, n_fft=2048, hop_length=512, n_mels=128))
+        add("forward", "tempogram", "librosa",
+            lambda: librosa.feature.tempogram(y=y, sr=SR, hop_length=512, win_length=384))
 
     if torch is not None and try_import("torchaudio"):
         import torchaudio.transforms as T
@@ -233,34 +269,6 @@ def build_cases(y, sec, torch, device, gpu, resident):
                     lambda td=td: (td(ytd), sync()),
                 )
 
-    if torch is not None and try_import("nnAudio"):
-        from nnAudio import features
-
-        def nn_make():
-            return {
-                "mel_spectrogram": features.MelSpectrogram(
-                    sr=SR,
-                    n_fft=2048,
-                    hop_length=512,
-                    n_mels=128,
-                    pad_mode="constant",
-                    verbose=False,
-                ),
-                "cqt": features.CQT(
-                    sr=SR, hop_length=512, n_bins=84, pad_mode="constant", verbose=False
-                ),
-            }
-
-        yt1 = torch.from_numpy(y)[None]
-        for name, t in nn_make().items():
-            add("forward", name, "nnAudio cpu", lambda t=t: t(yt1))
-        if device:
-            sync = syncer(torch, device)
-            ytd1 = yt1.to(device)
-            for name, t in nn_make().items():
-                td = t.to(device)
-                add("forward", name, f"nnAudio {device} resident", lambda td=td: (td(ytd1), sync()))
-
     if torch is None:
         return cases
 
@@ -287,16 +295,16 @@ def build_cases(y, sec, torch, device, gpu, resident):
         ("chroma", chroma_t),
     ]
 
-    for name, t in train:
+    for name, t in train + feat:
         add("training step", name, "auvux cpu", step(t, yt, no_sync))
     if gpu:
-        for name, t in train:
+        for name, t in train + feat:
             staged = step(lambda x, t=t: t(x, backend="gpu"), yt, no_sync)
             add("training step", name, f"auvux {gpu} staged", staged)
     if resident:
         sync = syncer(torch, device)
         ytd = torch.from_numpy(y).to(device)
-        for name, t in train:
+        for name, t in train + feat:
             add(
                 "training step",
                 name,
@@ -384,25 +392,13 @@ def build_cases(y, sec, torch, device, gpu, resident):
                 step(lambda x: inv_td(x, length=n_out), Ztd_ta, sync),
             )
 
-    if try_import("nnAudio"):
-        from nnAudio import features
+    return cases
 
-        ncqt = features.CQT(sr=SR, hop_length=512, n_bins=84, pad_mode="constant", verbose=False)
-        add("training step", "cqt", "nnAudio cpu", step(lambda x: ncqt(x[None]), yt, no_sync))
-        if device:
-            sync = syncer(torch, device)
-            ncqt_d = features.CQT(
-                sr=SR, hop_length=512, n_bins=84, pad_mode="constant", verbose=False
-            ).to(device)
-            ytd = torch.from_numpy(y).to(device)
-            add(
-                "training step",
-                "cqt",
-                f"nnAudio {device} resident",
-                step(lambda x: ncqt_d(x[None]), ytd, sync),
-            )
 
-    return cases
+# Transforms compared against librosa for speed only: their definitions match
+# librosa's values but not its envelope lag-shift convention, so a relerr "diff"
+# would read as a (misleading) failure rather than a correctness check.
+_NO_DIFF = frozenset({"onset_strength", "tempogram"})
 
 
 def run_cases(cases, repeats, warmup):
@@ -440,7 +436,10 @@ def run_cases(cases, repeats, warmup):
                     ts.append(time.perf_counter() - t0)
                     if rep == repeats - 1:
                         outs[id(c)] = r
-        ref = next((to_numpy(outs[id(c)]) for c, _, _ in live if c.impl == "librosa"), None)
+        grp_tr = live[0][0].transform if live else ""
+        ref = None
+        if grp_tr not in _NO_DIFF:
+            ref = next((to_numpy(outs[id(c)]) for c, _, _ in live if c.impl == "librosa"), None)
         if ref is not None:
             for c, _, _ in live:
                 if c.impl == "librosa":
@@ -466,7 +465,7 @@ def to_numpy(out):
     if hasattr(out, "detach"):
         out = out.detach().cpu().numpy()
     out = np.asarray(out)
-    while out.ndim > 2 and out.shape[0] == 1:  # nnAudio keeps a batch dim
+    while out.ndim > 2 and out.shape[0] == 1:  # drop a leading singleton batch dim
         out = out[0]
     return out
 
@@ -521,7 +520,7 @@ def print_table(title, col_names, rows):
         print(line(r))
 
 
-_LIB_ORDER = ("auvux", "librosa", "torchaudio", "nnAudio")
+_LIB_ORDER = ("auvux", "librosa", "torchaudio")
 _PATH_ORDER = ("cpu", "gpu staged", "gpu resident")
 
 
@@ -629,7 +628,7 @@ def main():
 
     torch, device, gpu, resident = detect_env()
     gpu_s = f"{gpu} ({dsp.backend_info(gpu)['device_name']})" if gpu else "none"
-    rivals = ("librosa", "torchaudio", "nnAudio")
+    rivals = ("librosa", "torchaudio")
     libs = ", ".join(f"{n} {version_of(n)}" for n in rivals if try_import(n))
     print(f"auvux-dsp {dsp.__version__} benchmark")
     threads = dsp.get_num_threads() or "auto"

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/pyproject.toml

@@ -20,13 +20,21 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
+# Full parity/grad suite: librosa is the forward oracle, torch the autograd
+# oracle. Used in dev and native-arch CI.
 test = ["pytest", "librosa", "soxr", "torch"]
+# Lean smoke test for the per-wheel cibuildwheel run: just confirm the compiled
+# extension loads and the native paths run. The librosa/torch oracle tests
+# importorskip and skip cleanly without those (heavy, and unbuildable on some
+# wheel targets — numba/llvmlite and torch have no x86_64-macOS wheels).
+wheeltest = ["pytest"]
 
 [tool.scikit-build]
 minimum-version = "build-system.requires"
 wheel.packages = ["python/auvux"]
 cmake.version = ">=3.18"
 build-dir = "build/{wheel_tag}"
+sdist.exclude = [".env", ".env.*", "*.env", ".pypirc", "*.pem", "wheelhouse", "dist"]
 
 [tool.scikit-build.metadata.version]
 provider = "scikit_build_core.metadata.regex"
@@ -35,7 +43,7 @@ input = "python/auvux/dsp/_version.py"
 [tool.cibuildwheel]
 build = "cp310-* cp311-* cp312-* cp313-*"
 skip = "*-musllinux* *_i686 *-win32"
-test-requires = "pytest"
+test-extras = ["wheeltest"]
 test-command = "pytest {project}/tests -q"
 
 [tool.cibuildwheel.macos]

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/python/auvux/dsp/__init__.py

@@ -26,22 +26,50 @@ from auvux.dsp._functional import (
     istft,
     mel_spectrogram,
     mfcc,
+    onset_strength,
+    rms,
+    spectral_bandwidth,
+    spectral_centroid,
+    spectral_flatness,
+    spectral_flux,
+    spectral_kurtosis,
+    spectral_moments,
+    spectral_skewness,
     stft,
+    tempogram,
     vqt,
 )
 from auvux.dsp._transform import Transform
-from auvux.dsp._transforms import CQT, ISTFT, MFCC, STFT, VQT, Chroma, MelSpectrogram
+from auvux.dsp._transforms import (
+    CQT,
+    ISTFT,
+    MFCC,
+    RMS,
+    STFT,
+    VQT,
+    Chroma,
+    MelSpectrogram,
+    OnsetStrength,
+    SpectralFlux,
+    SpectralMoments,
+    Tempogram,
+)
 from auvux.dsp._version import __version__
 
 __all__ = [
     "CQT",
     "ISTFT",
     "MFCC",
+    "RMS",
     "STFT",
     "VQT",
     "BackendError",
     "Chroma",
     "MelSpectrogram",
+    "OnsetStrength",
+    "SpectralFlux",
+    "SpectralMoments",
+    "Tempogram",
     "Transform",
     "__version__",
     "amplitude_to_db",
@@ -65,8 +93,18 @@ __all__ = [
     "mel_spectrogram",
     "mel_to_hz",
     "mfcc",
+    "onset_strength",
     "power_to_db",
+    "rms",
     "set_num_threads",
+    "spectral_bandwidth",
+    "spectral_centroid",
+    "spectral_flatness",
+    "spectral_flux",
+    "spectral_kurtosis",
+    "spectral_moments",
+    "spectral_skewness",
     "stft",
+    "tempogram",
     "vqt",
 ]

{auvux_dsp-0.1.0.dev0 → auvux_dsp-0.1.0.dev2}/python/auvux/dsp/_dispatch.py

@@ -17,7 +17,7 @@ class BackendError(RuntimeError):
 # Refuse to run against a compiled core older than these Python sources;
 # editable installs do not rebuild the extension on source changes.
 # Keep in lockstep with kAbiVersion in src/bindings/abi.hpp.
-_EXPECTED_ABI = 3
+_EXPECTED_ABI = 12
 if getattr(_native, "abi_version", lambda: 0)() != _EXPECTED_ABI:
     raise ImportError(
         "auvux.dsp's compiled core is out of date with its Python sources; "
@@ -102,7 +102,10 @@ def backend_info(name: str) -> dict[str, Any]:
             "available": True,
             "device_name": "",
             "supports": {
-                "transforms": ["stft", "istft", "mel_spectrogram", "mfcc", "cqt", "vqt", "chroma"],
+                "transforms": [
+                    "stft", "istft", "mel_spectrogram", "mfcc", "cqt", "vqt", "chroma",
+                    "spectral_moments", "rms", "spectral_flux", "onset_strength", "tempogram",
+                ],
                 "n_fft": "any power of two >= 32",
             },
         }
@@ -115,7 +118,10 @@ def backend_info(name: str) -> dict[str, Any]:
         "supports": {
             # MFCC rides on the GPU mel; istft is resident-only (device
             # tensors in, device tensors out).
-            "transforms": ["stft", "istft", "mel_spectrogram", "mfcc", "cqt", "vqt", "chroma"],
+            "transforms": [
+                "stft", "istft", "mel_spectrogram", "mfcc", "cqt", "vqt", "chroma", "rms",
+                "spectral_moments", "spectral_flux", "onset_strength", "tempogram",
+            ],
             "n_fft": "powers of two in [64, 4096]",
             "cqt": "hop_length divisible by 2**(n_octaves - 1)",
         },
@@ -198,12 +204,22 @@ def rewrap(out: np.ndarray, lead: tuple[int, ...], mod: str, device: Any = None)
     return out
 
 
-def _stream_for(kind: str) -> int:
-    if kind == "cuda":
+def _stream_for(t: Any, kind: str) -> int:
+    if kind != "cuda":
+        return 0
+    if array_module(t) == "torch":
         import torch
 
-        return int(torch.cuda.current_stream().cuda_stream) or 1
-    return 0
+        return int(torch.cuda.current_stream(getattr(t, "device", None)).cuda_stream) or 1
+    import cupy
+
+    return int(cupy.cuda.get_current_stream().ptr) or 1
+
+
+def _dlpack_capsule(t: Any, kind: str, stream: int) -> Any:
+    if kind == "cuda":
+        return t.__dlpack__(stream=stream)
+    return t.__dlpack__()
 
 
 def _frame_major(strides: tuple[int, ...], shape: tuple[int, ...], elem: int) -> bool:
@@ -225,13 +241,14 @@ def export_dlpack(
     cplx: bool = False,
     sync: bool = True,
     frame_major: bool = False,
+    stream: int | None = None,
 ) -> Any:
-    """GPU array (..., tail) -> (capsule, lead shape, keepalive) or None when
-    the array library cannot export DLPack here (then callers stage to host).
-    sync=False skips the device sync when an immediately preceding export
-    already drained the queue. frame_major=True normalizes a (..., bins,
-    frames) spectrum to the STFT/ISTFT engines' frame-major contract; other
-    engines (CQT complex included) take compact tensors."""
+    """GPU array (..., tail) -> (capsule, lead shape, keepalive, stream) or
+    None when the array library cannot export DLPack here (then callers stage
+    to host). sync=False skips the device sync when an immediately preceding
+    export already drained the queue. frame_major=True normalizes a (...,
+    bins, frames) spectrum to the STFT/ISTFT engines' frame-major contract;
+    other engines (CQT complex included) take compact tensors."""
     m = array_module(t)
     lead = tuple(t.shape[:-trailing])
     tail = tuple(t.shape[-trailing:])
@@ -252,8 +269,9 @@ def export_dlpack(
             t2 = t2.contiguous()
         if kind == "metal" and sync:
             torch.mps.synchronize()
+        stream = _stream_for(t2, kind) if stream is None else stream
         try:
-            return t2.__dlpack__(), lead, t2
+            return _dlpack_capsule(t2, kind, stream), lead, t2, stream
         except Exception:
             return None
     import cupy
@@ -265,8 +283,9 @@ def export_dlpack(
             t2 = cupy.ascontiguousarray(t2.transpose(0, 2, 1)).transpose(0, 2, 1)
     else:
         t2 = cupy.ascontiguousarray(t2)
+    stream = _stream_for(t2, kind) if stream is None else stream
     try:
-        return t2.__dlpack__(), lead, t2
+        return _dlpack_capsule(t2, kind, stream), lead, t2, stream
     except Exception:
         return None
 
@@ -289,8 +308,8 @@ def resident_istft(engine: Any, Z: Any, length: int, kind: str) -> Any:
     exp = export_dlpack(Z, kind, trailing=2, cplx=True, frame_major=True)
     if exp is None:
         return None
-    cap, lead, _keep = exp
-    out_cap = engine.forward_dlpack(cap, length, _stream_for(kind))
+    cap, lead, _keep, stream = exp
+    out_cap = engine.forward_dlpack(cap, length, stream)
     out = import_dlpack(out_cap, array_module(Z))
     return out.reshape(lead + tuple(out.shape[1:]))
 
@@ -300,10 +319,10 @@ def resident_forward(engine: Any, y: Any, output: str | None, kind: str) -> Any:
     exp = export_dlpack(y, kind, trailing=1)
     if exp is None:
         return None
-    cap, lead, _keep = exp
+    cap, lead, _keep, stream = exp
     if output is None:
-        out_cap = engine.forward_dlpack(cap, _stream_for(kind))
+        out_cap = engine.forward_dlpack(cap, stream)
     else:
-        out_cap = engine.forward_dlpack(cap, output, _stream_for(kind))
+        out_cap = engine.forward_dlpack(cap, output, stream)
     out = import_dlpack(out_cap, array_module(y))
     return out.reshape(lead + tuple(out.shape[1:]))