TT-Wavelet

TT-Wavelet implements Discrete Wavelet Transforms (DWT) on Tenstorrent accelerators through the Lifting Wavelet Transform (LWT). It provides one-level forward and inverse 1D and 2D transforms for TTNN tensors through the ttwt Python package.

Wavelet FIR filter banks are represented as sequences of lifting steps derived from Lifting Factorization of PyWavelets filter banks.

TT-Wavelet consumes the resulting decompositions. Its runtime does not embed that project's Python/SageMath factorization pipeline.

Features

• One-level 1D DWT / IDWT and 2D DWT / IDWT on device.
• 106 discrete wavelet schemes from the PyWavelets catalog.
• Eight boundary modes: zero, constant, symmetric, reflect, periodic, smooth, antisymmetric, and antireflect.
• Ordinary TTNN tensors as inputs and outputs with direct MemoryConfig interoperability.
• Package-local Wavelet JIT resources. TTNN / Metalium remain external dependencies.

The project catalog covers all 106 discrete wavelets from PyWavelets and all eight supported boundary modes. Numerical compatibility with PyWavelets has been validated for the supported one-level 1D and 2D DWT/IDWT operations, excluding dmey as described below.

Compatibility refers to numerical transform behavior, wavelet definitions, and boundary handling, not compatibility with the full PyWavelets Python API. periodic corresponds to periodic boundary extension and is distinct from PyWavelets' periodization, which is not supported.

dmey is excluded from numerical compatibility claims. Its lifting factorization has a large residual. See the factorization results. PyWavelets also omits dmey from its multilevel accuracy tests because its accuracy is very low.

Requirements

• Ubuntu 22.04 LTS, x86_64, CPython 3.10.
• ttnn==0.79.0.
• SFPI 7.78.0 [935], available from the official compiler release.

The official TTNN pip wheel does not bundle SFPI. A matching compiler installation must be available to TTNN at runtime, under ttnn/runtime/sfpi or /opt/tenstorrent/sfpi. TT-Wavelet does not download or bundle it.

Installation

PyPI

pip install tt-wavelet

The package declares ttnn==0.79.0 as a dependency. Pip installs that runtime provider. Driver, firmware, and SFPI provisioning remain separate prerequisites.

Development

The SUBMODULE build uses the pinned TT-Metal / TTNN v0.79.0 source baseline. Linux build prerequisites include clang-20, the GCC 12 C++ development toolchain, Ninja, pkg-config, and the libhwloc-dev, libnuma-dev, and libcapstone-dev packages.

git clone --recurse-submodules https://github.com/ke1rro/tt-wavelet.git
cd tt-wavelet

Build and install a local wheel:

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install build wheel cmake==4.2.3

CC=clang-20 \
CXX=clang++-20 \
CMAKE_BUILD_PARALLEL_LEVEL=2 \
python -m build --wheel \
  -Ccmake.define.TT_WAVELET_TT_METAL_MODE=SUBMODULE \
  -Cbuild-dir=build

python -m pip install dist/*.whl

Quick start

The example also requires PyTorch. The 1D API returns stick-padded tensors. Only the requested logical length of the reconstruction is valid.

import torch
import ttnn
import ttwt

device = ttnn.open_mesh_device(
    mesh_shape=ttnn.MeshShape(1, 1), physical_device_ids=[0]
)
try:
    signal = torch.arange(64, dtype=torch.float32)
    x = ttnn.from_torch(
        signal, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=device
    )
    cA, cD = ttwt.dwt(x, "db4")
    reconstructed = ttwt.idwt(cA, cD, "db4", original_length=signal.numel())
    result = ttnn.to_torch(reconstructed).flatten()[: signal.numel()]

    # 2D inputs use tile layout. The inverse takes the original logical shape.
    image = torch.arange(32 * 32, dtype=torch.float32).reshape(32, 32)
    x2d = ttnn.from_torch(
        image, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device
    )
    ll, lh, hl, hh = ttwt.dwt_2d(x2d, "db4")
    reconstructed = ttwt.idwt_2d(ll, lh, hl, hh, "db4", output_shape=image.shape)
finally:
    ttnn.close_mesh_device(device)

API

• ttwt.dwt_coeff_len
• ttwt.dwt
• ttwt.idwt
• ttwt.dwt_2d
• ttwt.idwt_2d

Hardware support

TT-Wavelet has been validated across all currently supported Wormhole and Blackhole SKUs targeted by TTNN 0.79.0.

Acknowledgements

TT-Wavelet originated from Tenstorrent bounty #40494. We thank Tenstorrent for sponsoring the work and providing the hardware used during development and validation.

References

• Lifting Factorization of PyWavelets filter banks.
• PyWavelets.
• PyWavelets: A Python package for wavelet analysis, Journal of Open Source Software.
• Tenstorrent TT-Metal / TTNN.
• Original LWT / ILWT bounty.

Citation

Please cite this software using CITATION.cff. TT-Wavelet authors: Nikita Lenyk and Andrii Kryvyi. The references above credit the separate factorization and PyWavelets work.

License

TT-Wavelet is licensed under the Apache License 2.0. Third-party and derived material retains its applicable licenses and notices. See licenses/ and NOTICE.