jaxwavelets

Execute JAX-native discrete, stationary, and continuous wavelet transforms on input arrays.

5|Updated Apr 14, 2026
One-click install
npx skills add https://github.com/handley-lab/jaxwavelets --skill jaxwavelets
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: jaxwavelets
Source: https://github.com/handley-lab/jaxwavelets/tree/main
Command: npx skills add https://github.com/handley-lab/jaxwavelets --skill jaxwavelets

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

This Skill helps you compute discrete, stationary, and continuous wavelet transforms inside JAX so you can differentiate through them, JIT-compile them, and run them efficiently on CPU/GPU.

Core Features & Use Cases

  • Discrete wavelet transforms (DWT): compute 1D natively (and extend to nD and 2D) with matching reconstruction via inverse transforms.
  • Stationary wavelet transforms (SWT): perform shift-invariant (undecimated) decompositions and reconstruct the original signal.
  • Continuous wavelet transforms (CWT): use a two-phase workflow (prepare then apply) to keep the core transform JIT-friendly.
  • Multiresolution analysis and wavelet packets: build interpretable multi-scale components and alternative decompositions.
  • Thresholding utilities: apply soft/hard/garrote/firm thresholding to wavelet coefficients.
  • JAX-first design: outputs are JAX pytrees for composition with jax.vmap, jax.jit, and jax.grad.

Quick Start

Use the jaxwavelets Skill to compute a 2D wavelet decomposition for a JAX array x with wavelet 'db4' at level 3.

Frequently Asked Questions about jaxwavelets

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I compute differentiable wavelet transforms in JAX?

You can compute differentiable wavelet transforms in JAX using native discrete, stationary, and continuous wavelet algorithms that output JAX pytrees, enabling seamless gradient computation through jax.grad.

Can I run JAX wavelet transforms on GPU and JIT-compile them?

Yes, JAX wavelet transforms are JIT-compatible and run efficiently on GPU, with the continuous wavelet transform using a two-phase prepare and apply kernel design to maintain JIT-friendliness.

Does this JAX wavelet library support a PyWavelets-matched API?

Yes, it provides a PyWavelets-matched API with single-example functional semantics, allowing users familiar with PyWavelets to perform 1D, nD, and 2D discrete wavelet transforms natively in JAX.

How do I apply thresholding to wavelet coefficients in JAX?

You can apply soft, hard, garrote, or firm thresholding directly to wavelet coefficients to perform signal denoising and multi-scale analysis within the JAX ecosystem.

What is the best way to perform shift-invariant wavelet decomposition in JAX?

Stationary wavelet transforms perform shift-invariant or undecimated decompositions in JAX, reconstructing the original signal while keeping outputs as pytree coefficients for batching and gradients.

When should I use continuous vs discrete wavelet transforms in JAX?

Use discrete wavelet transforms for standard multi-scale decomposition and reconstruction, and continuous wavelet transforms when you need a two-phase prepare and apply workflow that remains JIT-friendly.