Metadata-Version: 2.4
Name: specux
Version: 0.2.0.dev0
Summary: Differentiable audio DSP for Python: fast, fused kernels on GPU and CPU
Author-email: Peter Kiers <pkiers.1983@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://specux.com
Project-URL: Documentation, https://specux.com
Project-URL: Repository, https://github.com/auvux/specux
Project-URL: Issues, https://github.com/auvux/specux/issues
Keywords: stft,istft,spectrogram,mel,fft,dsp,audio,cuda,metal
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: C++
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Operating System :: MacOS
Classifier: Topic :: Multimedia :: Sound/Audio :: Analysis
Classifier: Topic :: Scientific/Engineering
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: numpy>=1.22
Provides-Extra: torch
Requires-Dist: torch>=2.4; extra == "torch"
Provides-Extra: jax
Requires-Dist: jax>=0.5; extra == "jax"
Provides-Extra: cuda
Requires-Dist: nvidia-cuda-nvrtc-cu12; extra == "cuda"
Requires-Dist: nvidia-cuda-runtime-cu12; extra == "cuda"
Provides-Extra: cuda11
Requires-Dist: nvidia-cuda-nvrtc-cu11; extra == "cuda11"
Requires-Dist: nvidia-cuda-runtime-cu11; extra == "cuda11"
Provides-Extra: cuda12
Requires-Dist: nvidia-cuda-nvrtc-cu12; extra == "cuda12"
Requires-Dist: nvidia-cuda-runtime-cu12; extra == "cuda12"
Provides-Extra: cuda13
Requires-Dist: nvidia-cuda-nvrtc<14,>=13; extra == "cuda13"
Requires-Dist: nvidia-cuda-runtime<14,>=13; extra == "cuda13"
Provides-Extra: rocm
Requires-Dist: rocm-sdk-core; extra == "rocm"
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: torch; extra == "test"
Provides-Extra: wheeltest
Requires-Dist: pytest; extra == "wheeltest"
Requires-Dist: jax>=0.5; extra == "wheeltest"
Dynamic: license-file

# SpecuX

[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)

Differentiable audio DSP for Python: fast, fused kernels on GPU and CPU.

SpecuX provides FFTs and convolution, spectral transforms and reconstruction,
audio features, filtering, and optional audio I/O through one functional API
for NumPy, Torch, CuPy, and JAX. Results use the input array library. Resident
arrays remain on their device.

Documentation: <https://specux.com>

```python
import specux
import torch

x = torch.randn(8, 32768, device="cuda", requires_grad=True)
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")
S.mean().backward()

assert S.device == x.device
assert x.grad is not None
```

## Execution model

Array libraries, compute engines, graph bridges, and memory transports are
separate concerns:

| Layer | Implementations | Responsibility |
|---|---|---|
| Array library | NumPy, Torch, CuPy, JAX | dtype, allocation, and device discovery |
| Compute engine | CPU, CUDA, Metal, ROCm | native kernels, plans, pools, and launch policy |
| Graph bridge | eager, Torch custom ops, XLA FFI | tracing, fake shapes, autograd, and VJPs |
| Transport | host, staged, device pointer, MTLBuffer, XLA buffer | moving or borrowing storage |

`backend=None` follows the input placement. An explicit `backend=` may stage a
NumPy array through another engine. Resident Torch, CuPy, and JAX arrays are
not redirected through the host. Move the array explicitly to change devices.
XLA controls JAX placement.

ROCm is represented as a separate engine rather than treated as CUDA: the
`specux._rocm` provider builds the same codegen kernels through HIP and hipRTC.
It serves resident PyTorch-ROCm and CuPy-ROCm arrays (whose device type is
still reported as `cuda`) by pointer, and stages NumPy through `backend="rocm"`.
PyTorch-ROCm is disambiguated from NVIDIA by `torch.version.hip`, so a HIP
tensor is never routed to the CUDA engine. JAX-on-ROCm runs through the same
XLA FFI bridge under the `ROCM` platform, so a resident JAX array stays on the
GPU across `jit`, `vmap`, and gradients.

## Features

- **FFT and convolution**: `fft`, `ifft`, `rfft`, and `irfft` at power-of-two,
  smooth, and prime lengths. Direct and FFT convolution support gradients for
  both operands.
- **Transforms and reconstruction**: `stft`, `istft`, `cqt`, `vqt`, `icqt`,
  and `griffinlim`, with complex, magnitude, power, and dB output modes.
- **Features** (`specux.feature`): mel spectrograms, MFCC, LFCC, chroma,
  RMS energy, onset strength, Tonnetz, spectral descriptors, deltas, inverse
  mel projection, and filterbank builders.
- **Filtering** (`specux.filters`): `lfilter`, `filtfilt`, differentiable
  biquad designs, `preemphasis`, and `deemphasis`. Signal and coefficient
  gradients are supported for `lfilter`.
- **Audio I/O** (`specux.audio`, optional): WAV, FLAC, MP3, OGG, and MP4
  decoding and encoding; frame-accurate slicing; batch and streaming I/O;
  resampling; loudness, true-peak, and loudness-range metering; and metadata.
- **Torch integration**: resident eager kernels, `torch.library` custom ops,
  analytic adjoints, autocast policies, fake kernels, and `torch.compile`
  support. `specux.transforms` provides `nn.Module` wrappers.
- **CuPy integration**: resident CUDA kernels use the array's device pointer
  and current stream without host staging.
- **JAX integration**: typed XLA FFI on CPU, CUDA, and ROCm, including
  `jax.jit`, gradients, and batching. Elementwise dB operations remain fusible
  JAX expressions. JAX Metal is not supported.

Float16 uses float32 compute by default. Eligible GPU transforms may use native
half compute with `compute="half"`. Metal computes in float32 and does not
support float64. JAX has no complex32, so half inputs widen to float32.

## Install

```bash
pip install specux            # NumPy runtime; platform wheels include native engines
pip install specux[torch]     # Torch autograd and compile integration
pip install specux[jax]       # JAX runtime for the typed XLA bridge
pip install specux[cuda]      # Recommended CUDA runtime, currently CUDA 12

pip install specux[cuda11]    # Explicit CUDA 11 runtime
pip install specux[cuda12]    # Explicit CUDA 12 runtime
pip install specux[cuda13]    # Explicit CUDA 13 runtime
pip install specux[rocm]      # AMD HIP runtime and hipRTC (rocm-sdk wheels)
```

SpecuX uses one CUDA engine for CUDA 11, 12, and 13. The optional extras only
install a matching NVRTC compiler and runtime headers. Use `cuda` for the
recommended version or a numbered extra to select one explicitly. `cuda`
remains an alias for `cuda12` throughout the SpecuX 0.2 release line. The
NVIDIA driver remains a system dependency.

No vendor library is linked: the CUDA driver and NVRTC, and the HIP runtime and
hipRTC, are all resolved when an engine is first used. One wheel therefore
carries the CPU, CUDA and ROCm engines, installs on a machine with none of
them, and binds whichever are present - including both GPU vendors at once. The
ROCm engine spans ROCm 6 and 7 from the same binary. AMD publishes the rocm-sdk
wheels on its own index rather than PyPI, so `specux[rocm]` needs
`--index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/`; a system ROCm
install needs nothing extra.

```python
specux.available("cuda")   # engine built in AND its runtime installed here
specux.devices()           # [Device(rocm:0 'AMD Radeon RX 7600S' gfx1102), ...]
```

Torch, CuPy, and JAX integrations are optional. Install the CuPy package that
matches your CUDA environment. SpecuX does not choose or install a CuPy
runtime. A base installation depends only on NumPy. Release wheels include the
XLA bridge but do not install JAX.

### Source builds

```bash
pip install -e .
```

An isolated build installs JAX in its temporary build environment to obtain the
typed FFI headers. This does not add JAX to the installed runtime requirements.
With `--no-build-isolation`, install JAX in the build environment before using
`SPECUX_JAX=1`.

Useful build controls:

- `SPECUX_JAX=0` omits the XLA bridge; `SPECUX_JAX=1` requires it.
- `SPECUX_CUDA=0` / `SPECUX_ROCM=0` omit that GPU engine; `=1` requires it.
  Both default to `auto` (build it when its SDK is found). The XLA bridge
  follows: it carries handlers for whichever GPU engines were built, and CUDA
  and ROCm can coexist in one build - they register under different XLA
  platforms - so a single wheel can serve both.
- `SPECUX_CPU_ONLY=1` builds only the CPU engine.
- `SPECUX_CUDA_HOME=/path/to/cuda` selects a CUDA toolkit.
- `SPECUX_AUDIO=0` omits audio I/O; `SPECUX_AUDIO=1` requires FFmpeg headers.
- `scripts/build.ps1` loads the MSVC environment and builds in place on
  Windows.

macOS source builds include the CPU and Metal engines. The Metal engine uses
the vendored metal-cpp headers. Audio I/O requires FFmpeg development
libraries, with TagLib used for tags and cover art. The DSP package remains
usable when the audio extension is omitted.

## Usage

### NumPy and explicit staging

```python
import numpy as np
import specux

x = np.random.default_rng(0).standard_normal((8, 32768)).astype(np.float32)

S = specux.stft(x, n_fft=1024, hop_length=256, output="power")
y = specux.istft(S, n_fft=1024, hop_length=256, length=x.shape[-1])
C = specux.feature.mfcc(x, sr=16000, n_mfcc=20)

# NumPy uses CPU by default; an explicit engine stages the host array.
S_cuda = specux.stft(x, n_fft=1024, backend="cuda")
```

### Audio I/O

```python
import specux

y, sr = specux.audio.load("take.wav", sr=16000, mono=True)
lufs = specux.audio.loudness(y, sr)
y = specux.audio.normalize(y, mode="lufs", target_db=-14.0, sr=sr)
specux.audio.save("out.flac", y, sr)

for block in specux.audio.blocks("session.flac", 30 * sr, sr=sr, mono=True):
    M = specux.feature.melspectrogram(block, sr=sr)
```

### Torch compile and autocast

```python
import specux
import torch

x = torch.randn(8, 32768, device="cuda")

compiled = torch.compile(
    lambda v: specux.stft(v, n_fft=1024, hop_length=256, output="power"),
    fullgraph=True,
)

with torch.autocast("cuda", torch.float16):
    S = compiled(x)
```

### CuPy resident CUDA

```python
import cupy as cp
import specux

x = cp.random.standard_normal((8, 32768)).astype(cp.float32)
S = specux.stft(x, n_fft=1024, hop_length=256, output="power")

assert isinstance(S, cp.ndarray)
assert S.device.id == x.device.id
```

SpecuX launches on CuPy's current CUDA stream and keeps inputs, intermediates,
and outputs on the device.

### JAX JIT and gradients

```python
import jax
import jax.numpy as jnp
import specux

x = jnp.ones((8, 32768), dtype=jnp.float32)

def loss(v):
    return specux.stft(v, n_fft=1024, hop_length=256, output="power").mean()

value, grad = jax.jit(jax.value_and_grad(loss))(x)
```

Enable `jax_enable_x64` before creating arrays when using float64:

```python
jax.config.update("jax_enable_x64", True)
```

## Runtime controls

- `specux.deterministic(True)` selects reproducible overlap-add kernels and
  follows `torch.use_deterministic_algorithms` for Torch callers.
- `specux.set_autotune(True)` tunes new GPU configurations and persists launch
  wisdom.
- `specux.autotune(...)` explicitly tunes one operation and configuration.
- `specux.clear_cache()` releases clearable Python caches, native plans, and
  device tables.
- `specux.clear_wisdom()` removes persisted launch wisdom.

## Design

SpecuX fuses framing, FFT, and output work when supported by the engine and
transform shape. Multi-stage operations remain on the same device and stream,
without intermediate copies through Python or host memory.

CUDA and Metal kernels share operation definitions through the accelerator
code generator. The CPU engine implements the same operation families with
AVX2 and NEON SIMD paths. Immutable windows, filterbanks, bases, and FFT data
live in bounded native plans. SpecuX resolves those plans behind the functional
API.

Torch integration uses a Python custom-op bridge, so the native engines do not
depend on the libtorch ABI. CuPy passes resident buffers and the current stream
directly to the CUDA engine. JAX integration uses a separate XLA FFI module, so
JAX remains optional at runtime.

Native source layout and layering rules are documented in
the [native source guide](src/README.md).

## Development

```bash
python scripts/lint.py
python -m pytest tests
python bench/bench_matrix.py
```

## License

MIT. Vendored third-party components and their licenses are listed in
[NOTICE](NOTICE).
