torchwavelets 0.1.0__py3-none-any.whl

torchwavelets/scatter/network.py

@@ -0,0 +1,204 @@
+"""General (arbitrary-order) DTCWT scattering network.
+
+A scattering network computes locally shift-invariant feature maps by
+iterating wavelet bandpasses with the smooth modulus ``sqrt(|Wx|^2 + b^2) - b``
+and smoothing the intermediate invariants with the wavelet lowpass. With the
+DTCWT's near shift-invariance this gives a small, cheap, highly invariant
+front-end for classification and segmentation.
+
+Paths are sequences of increasing scales ``(j1 < j2 < ... < jm)`` with
+``m <= max_order``. For each path the network returns one invariant:
+``avg_pool2d(lowpass, 2)`` when the path can still be extended (the smoothed
+modulus feeding the next order) and the raw ``|bandpass|`` at the terminal
+paths. With ``(J=2, max_order=2)`` the output matches
+:class:`~torchwavelets.scatter.lowlevel.ScatLayerj2` exactly.
+
+Shift invariance: the raw level-1 modulus ``|W_1 x|`` is only *weakly* shift
+invariant (the single-tree level-1 design), so individual raw bandpass
+invariants drift with sub-scale shifts; the smoothed invariants (pooled
+lowpasses) and the order-0 lowpass are exactly invariant at dyadic integer
+shifts, and the invariance improves with scale.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..dtcwt._autocast import autocast_off
+from ..dtcwt.coeffs import biort as _biort
+from ..dtcwt.coeffs import qshift as _qshift
+from ..dtcwt.transform2d import (
+    _compute_dtype,
+    _dtcwt2d_impl,
+    _prep_filt,
+    _validate_names,
+)
+
+__all__ = ["DTCWTScattering", "dtcwt_scattering"]
+
+
+def _smooth_mag(highs: torch.Tensor, magbias: float) -> torch.Tensor:
+    """Smooth modulus ``sqrt(|z|^2 + b^2) - b`` of complex bandpasses."""
+    return torch.sqrt(highs.abs().pow(2) + magbias**2) - magbias
+
+
+def _resolve_filters(biort, qshift, dtype, device):
+    _validate_names(biort, qshift)
+    h0o, _, h1o, _ = _biort(biort)
+    h0a, h0b, _, _, h1a, h1b, _, _ = _qshift(qshift)
+    return {
+        "h0o": _prep_filt(h0o, dtype, device),
+        "h1o": _prep_filt(h1o, dtype, device),
+        "h0a": _prep_filt(h0a, dtype, device),
+        "h0b": _prep_filt(h0b, dtype, device),
+        "h1a": _prep_filt(h1a, dtype, device),
+        "h1b": _prep_filt(h1b, dtype, device),
+    }
+
+
+def _dtcwt_levels(u: torch.Tensor, f, levels: int):
+    """``levels``-level DTCWT of ``u``; returns ``(lowpass, bandpass)``."""
+    yl, yh = _dtcwt2d_impl(u, f, levels, [False] * levels, [False] * levels, 2)
+    return yl, yh[-1]
+
+
+def _scattering_impl(x, f, J, max_order, magbias):
+    """BFS path expansion; returns ``(s0, [(path, invariant), ...])``.
+
+    Invariant semantics (matching ScatLayerj2): the invariant of an
+    extendable path is the smoothed (avg-pooled) lowpass of the level-1
+    transform of its modulus; terminal paths return the raw modulus.
+    """
+    if J < 1:
+        raise ValueError("J must be >= 1 for a scattering network")
+    # order-0 invariant: the smoothed final lowpass
+    yl, _ = _dtcwt_levels(x, f, J)
+    s0 = F.avg_pool2d(yl, 2)
+
+    results = []
+    # frontier entries: (u, current_scale, path, pending_path)
+    # pending_path = the path whose smoothed invariant comes from the level-1
+    # transform of u (computed at this node).
+    frontier = [(x, 0, (), None)]
+    for _ in range(max_order):
+        next_frontier = []
+        for u, scale, path, pending in frontier:
+            # level-1 transform of u: needed for the pending invariant and as
+            # the first step of the j-loop below.
+            ll1, hh1 = _dtcwt_levels(u, f, 1)
+            u1 = _flatten_orientations(_smooth_mag(hh1, magbias))
+            if pending is not None:
+                results.append((pending, F.avg_pool2d(ll1, 2)))
+            for j in range(scale + 1, J + 1):
+                if j == scale + 1:
+                    ll, hh, u_next = ll1, hh1, u1
+                else:
+                    ll, hh = _dtcwt_levels(u, f, j - scale)
+                    u_next = _flatten_orientations(_smooth_mag(hh, magbias))
+                new_path = path + (j,)
+                if j < J and len(new_path) < max_order:
+                    next_frontier.append((u_next, j, new_path, new_path))
+                else:
+                    results.append((new_path, _upstream_order(u_next, len(new_path))))
+        frontier = next_frontier
+    results.sort(key=lambda pv: (len(pv[0]), pv[0]))
+    return s0, results
+
+
+def _flatten_orientations(u: torch.Tensor) -> torch.Tensor:
+    """``(N, C, 6, H, W)`` modulus -> ``(N, 6*C, H, W)``."""
+    b, ch, _, h, w = u.shape
+    return u.reshape(b, ch * 6, h, w)
+
+
+def _upstream_order(u: torch.Tensor, m: int) -> torch.Tensor:
+    """Rotate the orientation blocks to the reference channel order.
+
+    Internally the moduli propagate channel-major as
+    ``(C, o1st, o2nd, ..., om)``; the reference (``ScatLayerj2``) stacks the
+    first orientation block last: ``(C, o2nd, ..., om, o1st)``. For a
+    first-order modulus (``m == 1``) this is the identity — the reference
+    keeps the moduli channel-major ``(C, o1st)``.
+    """
+    b, c6, h, w = u.shape
+    c = c6 // 6**m
+    blocks = u.reshape(b, c, *([6] * m), h, w)
+    perm = (0, 1) + tuple(range(3, m + 2)) + (2, m + 2, m + 3)
+    return blocks.permute(perm).reshape(b, c6, h, w)
+
+
+def dtcwt_scattering(
+    x: torch.Tensor,
+    J: int = 3,
+    max_order: int = 2,
+    biort="near_sym_a",
+    qshift="qshift_a",
+    magbias: float = 1e-2,
+) -> tuple[torch.Tensor, list[tuple[tuple[int, ...], torch.Tensor]]]:
+    """Forward DTCWT scattering network.
+
+    Args:
+        x (torch.Tensor): ``(N, C, H, W)`` input.
+        J (int): number of wavelet scales.
+        max_order (int): maximum scattering order.
+        biort (str): filter bank name.
+        qshift (str): filter bank name.
+        magbias (float): smoothing bias of the modulus.
+
+    Returns:
+        tuple[torch.Tensor, list[tuple[tuple[int, ...], torch.Tensor]]]:
+        ``(s0, invariants)``: the smoothed order-0 lowpass
+        ``(N, C, H/2^J, W/2^J)`` and a list of ``(path, tensor)`` pairs, one
+        per path (orders 1..max_order, ordered by order then scale). The
+        invariant size depends only on the path's final scale ``jm``:
+        ``H/2^jm``.
+    """
+    if x.ndim != 4:
+        raise ValueError(f"dtcwt_scattering expects (N, C, H, W); got {tuple(x.shape)}")
+    if isinstance(J, bool) or not isinstance(J, int) or J < 1:
+        raise ValueError(f"J must be a positive int; got {J!r}")
+    if isinstance(max_order, bool) or not isinstance(max_order, int) or max_order < 1:
+        raise ValueError(f"max_order must be a positive int; got {max_order!r}")
+    dtype = _compute_dtype(x)
+    f = _resolve_filters(biort, qshift, dtype, x.device)
+    with autocast_off(x.device):
+        return _scattering_impl(x.to(dtype), f, J, max_order, float(magbias))
+
+
+class DTCWTScattering(nn.Module):
+    """DTCWT scattering network module (see :func:`dtcwt_scattering`)."""
+
+    def __init__(self, J=3, max_order=2, biort="near_sym_a", qshift="qshift_a", magbias=1e-2):
+        super().__init__()
+        if isinstance(J, bool) or not isinstance(J, int) or J < 1:
+            raise ValueError(f"J must be a positive int; got {J!r}")
+        if isinstance(max_order, bool) or not isinstance(max_order, int) or max_order < 1:
+            raise ValueError(f"max_order must be a positive int; got {max_order!r}")
+        self.J = J
+        self.max_order = max_order
+        self.biort = biort
+        self.qshift = qshift
+        self.magbias = float(magbias)
+        filts = _resolve_filters(biort, qshift, torch.float64, torch.device("cpu"))
+        for name, t in filts.items():
+            self.register_buffer(name, t)
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim != 4:
+            raise ValueError(f"DTCWTScattering expects (N, C, H, W); got {tuple(x.shape)}")
+        dtype = _compute_dtype(x)
+        f = {
+            n: getattr(self, n).to(dtype=dtype, device=x.device)
+            for n in ("h0o", "h1o", "h0a", "h0b", "h1a", "h1b")
+        }
+        with autocast_off(x.device):
+            return _scattering_impl(x.to(dtype), f, self.J, self.max_order, self.magbias)
+
+    def extra_repr(self) -> str:
+        return (
+            f"J={self.J}, max_order={self.max_order}, "
+            f"biort={self.biort!r}, qshift={self.qshift!r}, "
+            f"magbias={self.magbias}"
+        )

torchwavelets/swt/__init__.py

@@ -0,0 +1,14 @@
+"""Stationary (undecimated) wavelet transforms."""
+
+from .transform import (
+    ISWT1D,
+    ISWT2D,
+    SWT1D,
+    SWT2D,
+    iswt1d,
+    iswt2d,
+    swt1d,
+    swt2d,
+)
+
+__all__ = ["swt1d", "iswt1d", "swt2d", "iswt2d", "SWT1D", "ISWT1D", "SWT2D", "ISWT2D"]

torchwavelets/swt/lowlevel.py

@@ -0,0 +1,53 @@
+"""Low-level stationary (undecimated / a trous) wavelet filter banks.
+
+The SWT is the undecimated periodic wavelet transform: at level ``j`` the
+analysis filters are dilated (zero-stuffed) by ``2**(j-1)`` and applied with
+no decimation, so every subband has the same length as the input. The
+phase conventions match PyWavelets exactly (verified bit-for-bit by the
+parity tests).
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+__all__ = ["swt_axis", "dilate_filter"]
+
+
+def dilate_filter(h: torch.Tensor, fstep: int) -> torch.Tensor:
+    """Insert ``fstep - 1`` zeros between filter taps (a trous dilation)."""
+    L = h.numel()
+    e = torch.zeros(1, 1, L * fstep, 1, dtype=h.dtype, device=h.device)
+    e[0, 0, ::fstep, 0] = h.reshape(-1)
+    return e
+
+
+def swt_axis(x: torch.Tensor, h: torch.Tensor, dim: int, fstep: int) -> torch.Tensor:
+    """One stationary analysis filter along ``dim`` (no decimation).
+
+    Args:
+        x: ``(N, C, H, W)`` input.
+        h: reversed analysis filter, ``(1, 1, L, 1)`` (column orientation;
+            transposed internally for ``dim == 3``).
+        dim: spatial dimension (2 or 3).
+        fstep: dilation factor ``2**(level - 1)``.
+
+    Returns:
+        ``(N, C, ...)`` with the same spatial size as ``x``.
+    """
+    n = x.shape[dim]
+    L = h.numel()
+    Ld = L * fstep  # dilated filter length
+    e = dilate_filter(h, fstep)
+    if dim == 3:
+        e = e.transpose(2, 3).contiguous()
+    # pywt: out[i] = sum_m dec[m] x[(i + (L/2 - m)*fstep) mod n]; with the
+    # reversed-filter convention this is out[i] = sum_k h_rev[k]
+    # x[(i + (k - L/2 + 1)*fstep) mod n].
+    shift = fstep * (L // 2 - 1)
+    idx = (torch.arange(n + Ld - 1, device=x.device) - shift) % n
+    xp = x.index_select(dim, idx)
+    stride = (1, 1)
+    C = x.shape[1]
+    return F.conv2d(xp, e.repeat(C, 1, 1, 1), stride=stride, groups=C)

torchwavelets/swt/transform.py

@@ -0,0 +1,386 @@
+"""Stationary (undecimated) wavelet transform: functional and module API.
+
+The SWT (algorithm a trous) keeps every subband at the input resolution by
+dilating the analysis filters by ``2**(level-1)`` and skipping decimation.
+Coefficients follow the PyWavelets conventions exactly (coarsest first), and
+``iswt`` / ``iswt2`` are exact inverses (verified bit-for-bit against pywt
+by the parity tests).
+
+Inputs must be divisible by ``2**level`` along every transformed axis, as
+required by the stationary transform (PyWavelets raises the same error).
+"""
+
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+import torch.nn as nn
+
+from ..core.filters import resolve_2d
+from ..dwt import lowlevel as dwt_ll
+from ..dwt._validation import check_spatial
+from .lowlevel import swt_axis
+
+__all__ = ["swt1d", "iswt1d", "swt2d", "iswt2d", "SWT1D", "ISWT1D", "SWT2D", "ISWT2D"]
+
+
+def _check_level(level: int) -> None:
+    if isinstance(level, bool) or not isinstance(level, int) or level < 1:
+        raise ValueError(f"level must be a positive int; got {level!r}")
+
+
+def _check_divisible(n: int, level: int, name: str) -> None:
+    if n % (2**level) != 0:
+        raise ValueError(
+            f"{name} length {n} is not divisible by 2**level = {2**level}; "
+            "the stationary wavelet transform requires it (pad the input)"
+        )
+
+
+def _iswt_idwt(lo: torch.Tensor, hi: torch.Tensor, g0r, g1r) -> torch.Tensor:
+    """Single-level periodic inverse DWT on 3-D ``(N, C, L)`` coefficients."""
+    lo4 = lo.unsqueeze(2)
+    hi4 = hi.unsqueeze(2)
+    return dwt_ll.sfb1d(lo4, hi4, g0r, g1r, "periodization", 3).squeeze(2)
+
+
+def swt1d(
+    x: torch.Tensor, wave, level: int, norm: bool = False
+) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    """1-D stationary wavelet transform.
+
+    Args:
+        x (torch.Tensor): ``(N, C, L)`` input.
+        wave (str or pywt.Wavelet or tuple): wavelet name / bank (see
+            :func:`torchwavelets.core.filters.filter_bank`).
+        level (int): number of scales.
+        norm (bool): normalize so the coefficient energy equals the input
+            energy (orthogonal wavelets only, like PyWavelets).
+
+    Returns:
+        list[tuple[torch.Tensor, torch.Tensor]]: ``[(cA_n, cD_n), ...,
+        (cA_1, cD_1)]`` (coarsest first), every subband of shape ``(N, C, L)``.
+    """
+    _check_level(level)
+    if x.ndim != 3:
+        raise ValueError(f"swt1d expects (N, C, L); got {tuple(x.shape)}")
+    check_spatial(x.shape[-1:], ("L",), "swt1d")
+    _check_divisible(x.shape[-1], level, "signal")
+    h0c, h1c, _, _, g0c, g1c, _, _ = resolve_2d(wave, x.dtype, x.device)
+    if norm:
+        h0c, h1c = h0c / 2**0.5, h1c / 2**0.5
+    x4 = x.unsqueeze(-1)  # (N, C, L, 1)
+    ret = []
+    a = x4
+    for j in range(1, level + 1):
+        fstep = 2 ** (j - 1)
+        ca = swt_axis(a, h0c, 2, fstep)
+        cd = swt_axis(a, h1c, 2, fstep)
+        ret.append((ca.squeeze(-1), cd.squeeze(-1)))
+        a = ca
+    ret.reverse()
+    return ret
+
+
+def iswt1d(coeffs, wave, norm: bool = False) -> torch.Tensor:
+    """Inverse 1-D stationary wavelet transform (exact inverse of :func:`swt1d`)."""
+    for cA, cD in coeffs:
+        if cA.shape != cD.shape:
+            raise ValueError(
+                "iswt1d requires cA and cD of the same shape; got "
+                f"{tuple(cA.shape)} vs {tuple(cD.shape)}"
+            )
+    _, _, _, _, _, _, g0r, g1r = resolve_2d(wave, coeffs[0][0].dtype, coeffs[0][0].device)
+    if norm:
+        g0r, g1r = g0r * 2**0.5, g1r * 2**0.5
+    n_levels = len(coeffs)
+    output = coeffs[0][0].clone()
+    N = output.shape[-1]
+    for j in range(n_levels, 0, -1):
+        step = 2 ** (j - 1)
+        cD = coeffs[-j][1]
+        for first in range(step):
+            idx = torch.arange(first, N, step, device=output.device)
+            even = idx[0::2]
+            odd = idx[1::2]
+            x1 = _iswt_idwt(
+                output.index_select(-1, even),
+                cD.index_select(-1, even),
+                g0r,
+                g1r,
+            )
+            x2 = _iswt_idwt(
+                output.index_select(-1, odd),
+                cD.index_select(-1, odd),
+                g0r,
+                g1r,
+            )
+            x2 = torch.roll(x2, 1, dims=-1)
+            output[..., idx] = (x1 + x2) / 2
+    return output
+
+
+def swt2d(
+    x: torch.Tensor, wave, level: int, norm: bool = False
+) -> list[tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]:
+    """2-D stationary wavelet transform.
+
+    Args:
+        x (torch.Tensor): ``(N, C, H, W)`` input.
+        wave (str or pywt.Wavelet or tuple): as in :func:`swt1d`.
+        level (int): as in :func:`swt1d`.
+        norm (bool): as in :func:`swt1d`.
+
+    Returns:
+        list[tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]:
+        ``[(cA_j, (cH_j, cV_j, cD_j)), ...]`` (coarsest first), all subbands
+        of shape ``(N, C, H, W)``.
+    """
+    _check_level(level)
+    if x.ndim != 4:
+        raise ValueError(f"swt2d expects (N, C, H, W); got {tuple(x.shape)}")
+    check_spatial(x.shape[2:], ("H", "W"), "swt2d")
+    _check_divisible(x.shape[-2], level, "height")
+    _check_divisible(x.shape[-1], level, "width")
+    h0c, h1c, h0r, h1r, _, _, _, _ = resolve_2d(wave, x.dtype, x.device)
+    if norm:
+        h0c, h1c, h0r, h1r = h0c / 2**0.5, h1c / 2**0.5, h0r / 2**0.5, h1r / 2**0.5
+    ret = []
+    a = x
+    for j in range(1, level + 1):
+        fstep = 2 ** (j - 1)
+        # axis 0 (rows, dim 2), then axis 1 (cols, dim 3), like pywt.swt2
+        low = swt_axis(a, h0c, 2, fstep)
+        high = swt_axis(a, h1c, 2, fstep)
+        aa = swt_axis(low, h0r, 3, fstep)
+        ad = swt_axis(low, h1r, 3, fstep)
+        da = swt_axis(high, h0r, 3, fstep)
+        dd = swt_axis(high, h1r, 3, fstep)
+        ret.append((aa, (da, ad, dd)))
+        a = aa
+    ret.reverse()
+    return ret
+
+
+def iswt2d(coeffs, wave, norm: bool = False) -> torch.Tensor:
+    """Inverse 2-D stationary wavelet transform (exact inverse of :func:`swt2d`)."""
+    for cA, cs in coeffs:
+        if any(c.shape != cA.shape for c in cs):
+            raise ValueError(
+                "iswt2d requires all subbands of a level to share the shape "
+                f"{tuple(cA.shape)}; got {[tuple(c.shape) for c in cs]}"
+            )
+    _, _, _, _, g0c, g1c, g0r, g1r = resolve_2d(wave, coeffs[0][0].dtype, coeffs[0][0].device)
+    if norm:
+        g0c, g1c, g0r, g1r = g0c * 2**0.5, g1c * 2**0.5, g0r * 2**0.5, g1r * 2**0.5
+    n_levels = len(coeffs)
+    output = coeffs[0][0].clone()
+    H, W = output.shape[-2:]
+    for j in range(n_levels):
+        step = 2 ** (n_levels - j - 1)
+        cH, cV, cD = coeffs[j][1]
+        for fh in range(step):
+            for fw in range(step):
+                approx = output.clone()
+                out_slice = output[..., fh::step, fw::step]
+                out_slice.zero_()
+                for oh in (0, 1):
+                    for ow in (0, 1):
+                        h_eo = slice(fh + oh * step, H, 2 * step)
+                        w_eo = slice(fw + ow * step, W, 2 * step)
+                        ll = approx[..., h_eo, w_eo]
+                        lh = cH[..., h_eo, w_eo]
+                        hl = cV[..., h_eo, w_eo]
+                        hh = cD[..., h_eo, w_eo]
+                        lo = dwt_ll.sfb1d(ll, lh, g0c, g1c, "periodization", 2)
+                        hi = dwt_ll.sfb1d(hl, hh, g0c, g1c, "periodization", 2)
+                        xr = dwt_ll.sfb1d(lo, hi, g0r, g1r, "periodization", 3)
+                        if oh:
+                            xr = torch.roll(xr, 1, dims=-2)
+                        if ow:
+                            xr = torch.roll(xr, 1, dims=-1)
+                        out_slice += xr
+                out_slice /= 4
+    return output
+
+
+class SWT1D(nn.Module):
+    """1-D stationary wavelet transform module (see :func:`swt1d`)."""
+
+    def __init__(self, wave, level: int, norm: bool = False):
+        super().__init__()
+        _check_level(level)
+        self.wave = wave
+        self.level = level
+        self.norm = norm
+        h0c, h1c, _, _, _, _, _, _ = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        self.register_buffer("h0c", h0c)
+        self.register_buffer("h1c", h1c)
+
+    def forward(self, x: torch.Tensor):
+        h0c = cast(torch.Tensor, self.h0c).to(dtype=x.dtype, device=x.device)
+        h1c = cast(torch.Tensor, self.h1c).to(dtype=x.dtype, device=x.device)
+        if self.norm:
+            h0c, h1c = h0c / 2**0.5, h1c / 2**0.5
+        _check_level(self.level)
+        if x.ndim != 3:
+            raise ValueError(f"SWT1D expects (N, C, L); got {tuple(x.shape)}")
+        _check_divisible(x.shape[-1], self.level, "signal")
+        x4 = x.unsqueeze(-1)
+        ret = []
+        a = x4
+        for j in range(1, self.level + 1):
+            fstep = 2 ** (j - 1)
+            ca = swt_axis(a, h0c, 2, fstep)
+            cd = swt_axis(a, h1c, 2, fstep)
+            ret.append((ca.squeeze(-1), cd.squeeze(-1)))
+            a = ca
+        ret.reverse()
+        return ret
+
+    def extra_repr(self) -> str:
+        return f"wave={self.wave!r}, level={self.level}, norm={self.norm}"
+
+
+class ISWT1D(nn.Module):
+    """1-D inverse stationary wavelet transform module (see :func:`iswt1d`)."""
+
+    def __init__(self, wave, norm: bool = False):
+        super().__init__()
+        self.wave = wave
+        self.norm = norm
+        _, _, _, _, _, _, g0r, g1r = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        self.register_buffer("g0r", g0r)
+        self.register_buffer("g1r", g1r)
+
+    def forward(self, coeffs):
+        g0r = self.g0r.to(dtype=coeffs[0][0].dtype, device=coeffs[0][0].device)
+        g1r = self.g1r.to(dtype=coeffs[0][0].dtype, device=coeffs[0][0].device)
+        if self.norm:
+            g0r, g1r = g0r * 2**0.5, g1r * 2**0.5
+        n_levels = len(coeffs)
+        output = coeffs[0][0].clone()
+        N = output.shape[-1]
+        for j in range(n_levels, 0, -1):
+            step = 2 ** (j - 1)
+            cD = coeffs[-j][1]
+            for first in range(step):
+                idx = torch.arange(first, N, step, device=output.device)
+                even = idx[0::2]
+                odd = idx[1::2]
+                x1 = _iswt_idwt(
+                    output.index_select(-1, even),
+                    cD.index_select(-1, even),
+                    g0r,
+                    g1r,
+                )
+                x2 = _iswt_idwt(
+                    output.index_select(-1, odd),
+                    cD.index_select(-1, odd),
+                    g0r,
+                    g1r,
+                )
+                x2 = torch.roll(x2, 1, dims=-1)
+                output[..., idx] = (x1 + x2) / 2
+        return output
+
+    def extra_repr(self) -> str:
+        return f"wave={self.wave!r}, norm={self.norm}"
+
+
+class SWT2D(nn.Module):
+    """2-D stationary wavelet transform module (see :func:`swt2d`)."""
+
+    def __init__(self, wave, level: int, norm: bool = False):
+        super().__init__()
+        _check_level(level)
+        self.wave = wave
+        self.level = level
+        self.norm = norm
+        h0c, h1c, h0r, h1r, _, _, _, _ = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        for name, t in (("h0c", h0c), ("h1c", h1c), ("h0r", h0r), ("h1r", h1r)):
+            self.register_buffer(name, t)
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim != 4:
+            raise ValueError(f"SWT2D expects (N, C, H, W); got {tuple(x.shape)}")
+        _check_divisible(x.shape[-2], self.level, "height")
+        _check_divisible(x.shape[-1], self.level, "width")
+        f = {
+            n: getattr(self, n).to(dtype=x.dtype, device=x.device)
+            for n in ("h0c", "h1c", "h0r", "h1r")
+        }
+        if self.norm:
+            for n in f:
+                f[n] = f[n] / 2**0.5
+        ret = []
+        a = x
+        for j in range(1, self.level + 1):
+            fstep = 2 ** (j - 1)
+            low = swt_axis(a, f["h0c"], 2, fstep)
+            high = swt_axis(a, f["h1c"], 2, fstep)
+            aa = swt_axis(low, f["h0r"], 3, fstep)
+            ad = swt_axis(low, f["h1r"], 3, fstep)
+            da = swt_axis(high, f["h0r"], 3, fstep)
+            dd = swt_axis(high, f["h1r"], 3, fstep)
+            ret.append((aa, (da, ad, dd)))
+            a = aa
+        ret.reverse()
+        return ret
+
+    def extra_repr(self) -> str:
+        return f"wave={self.wave!r}, level={self.level}, norm={self.norm}"
+
+
+class ISWT2D(nn.Module):
+    """2-D inverse stationary wavelet transform module (see :func:`iswt2d`)."""
+
+    def __init__(self, wave, norm: bool = False):
+        super().__init__()
+        self.wave = wave
+        self.norm = norm
+        _, _, _, _, g0c, g1c, g0r, g1r = resolve_2d(wave, torch.float64, torch.device("cpu"))
+        for name, t in (("g0c", g0c), ("g1c", g1c), ("g0r", g0r), ("g1r", g1r)):
+            self.register_buffer(name, t)
+
+    def forward(self, coeffs):
+        f = {
+            n: getattr(self, n).to(dtype=coeffs[0][0].dtype, device=coeffs[0][0].device)
+            for n in ("g0c", "g1c", "g0r", "g1r")
+        }
+        if self.norm:
+            for n in f:
+                f[n] = f[n] * 2**0.5
+        n_levels = len(coeffs)
+        output = coeffs[0][0].clone()
+        H, W = output.shape[-2:]
+        for j in range(n_levels):
+            step = 2 ** (n_levels - j - 1)
+            cH, cV, cD = coeffs[j][1]
+            for fh in range(step):
+                for fw in range(step):
+                    approx = output.clone()
+                    out_slice = output[..., fh::step, fw::step]
+                    out_slice.zero_()
+                    for oh in (0, 1):
+                        for ow in (0, 1):
+                            h_eo = slice(fh + oh * step, H, 2 * step)
+                            w_eo = slice(fw + ow * step, W, 2 * step)
+                            ll = approx[..., h_eo, w_eo]
+                            lh = cH[..., h_eo, w_eo]
+                            hl = cV[..., h_eo, w_eo]
+                            hh = cD[..., h_eo, w_eo]
+                            lo = dwt_ll.sfb1d(ll, lh, f["g0c"], f["g1c"], "periodization", 2)
+                            hi = dwt_ll.sfb1d(hl, hh, f["g0c"], f["g1c"], "periodization", 2)
+                            xr = dwt_ll.sfb1d(lo, hi, f["g0r"], f["g1r"], "periodization", 3)
+                            if oh:
+                                xr = torch.roll(xr, 1, dims=-2)
+                            if ow:
+                                xr = torch.roll(xr, 1, dims=-1)
+                            out_slice += xr
+                    out_slice /= 4
+        return output
+
+    def extra_repr(self) -> str:
+        return f"wave={self.wave!r}, norm={self.norm}"