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.