code/models/tt_dit/layers/audio_ops.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC

# SPDX-License-Identifier: Apache-2.0

"""Audio conv primitives for LTX-2 audio decode (1D/2D convs via conv3d)."""

from __future__ import annotations

import math
from typing import Sequence

import torch
from loguru import logger

import ttnn

from ..layers.module import Module, Parameter
from ..parallel.config import AudioTCParallelConfig, AudioTParallelConfig, ParallelFactor
from ..parallel.manager import CCLManager
from ..utils.conv3d import _ntuple, aligned_channels, get_conv3d_config
from ..utils.tap_filter_configs import (
    applicable_formulations,
    format_formulation_row,
    slice_signature,
    tap_device_key,
    tap_formulation,
    tap_slice_config,
)
from ..utils.tensor import local_device_to_torch

# Per-mesh cache of constant zeros buffers, keyed by id(mesh_device).
_ZEROS_CACHE: dict = {}

CONV_SPLIT_MODES = ("off", "weight", "full")

# Default cap on conv3d's C_in_block for the H3 audio blocking table; the sweep that keeps it at
# 128 lives in `blockings_minimax_h3_audio`.
DEFAULT_MAX_C_IN_BLOCK = 128


def weights_variant(split_mode: str, max_c_in_block: int = DEFAULT_MAX_C_IN_BLOCK) -> str:
    """Cache-key suffix for the precision levers that change the prepared parameter set.

    ``split_mode`` decides whether the ``weight_lo`` residual parameters exist, so device-weight
    caches prepared under different settings hold different ``.tensorbin`` sets and are not
    interchangeable. ``max_c_in_block`` changes the prepared weight *bytes* with an unchanged file
    set (`prepare_conv3d_weight_state` blocks by ``C_in_block``), so it gets a term too.

    Only the all-fast configuration maps to ``""``; the H3 default yields ``"_split-full"`` --
    deliberately distinct from both the unsuffixed fast path and the retired tap_matmul era's
    ``"_split-full_tap1"``, so it can never collide with a stale cache from either. The caller
    must pass the same values it constructed its modules with.
    """
    suffix = "" if split_mode == "off" else f"_split-{split_mode}"
    if max_c_in_block != DEFAULT_MAX_C_IN_BLOCK:
        suffix += f"_cinb{max_c_in_block}"
    return suffix


def _split_operand(x: ttnn.Tensor) -> tuple[ttnn.Tensor, ttnn.Tensor]:
    """``(hi, lo)`` with ``hi = bf16(x)`` and ``lo = x - hi``, exactly, staying in the input's layout.

    Verified on ROW_MAJOR fp32: the typecast round trip reproduces torch's ``.bfloat16()`` bit-for-bit and
    ``hi + lo`` reconstructs ``x`` exactly, so this costs two elementwise ops and no tilize -- which
    matters, since the audio decode stage profiles at ~70 % layout ops.
    """
    hi = ttnn.typecast(ttnn.typecast(x, ttnn.bfloat16), ttnn.float32)
    return hi, ttnn.subtract(x, hi)


def conv3d_maybe_split(
    *,
    split_mode: str,
    input_tensor: ttnn.Tensor,
    weight_tensor: ttnn.Tensor,
    weight_lo_tensor: ttnn.Tensor | None,
    bias_tensor: ttnn.Tensor | None,
    **conv_kwargs,
) -> ttnn.Tensor:
    """``ttnn.experimental.conv3d``, optionally summed over operand-split terms.

    The fp32 conv3d multiply keeps only ~11 significand bits, and its relative error is flat in the
    reduction depth -- the operands are being truncated, not the sum drifting, so ``fp32_dest_acc_en``
    cannot help. Conv is linear in both arguments, so splitting an operand into ``hi = bf16(v)`` plus
    the exact residual ``lo = v - hi`` lets a second conv carry the dropped mantissa bits.
    ``split_mode="weight"`` splits the weight only (2 convs, measured 1.5x less error on ``conv_pre``);
    ``"full"`` splits both (3 convs -- the ``lo*lo`` term is negligible and omitted; 1.9x).

    ``bias`` is applied to exactly one term, since it is not a factor of the product being split.
    """
    if split_mode == "off":
        return ttnn.experimental.conv3d(
            input_tensor=input_tensor, weight_tensor=weight_tensor, bias_tensor=bias_tensor, **conv_kwargs
        )
    assert weight_lo_tensor is not None, f"split_mode={split_mode!r} needs a prepared weight residual"

    if split_mode == "weight":
        x_hi, terms = input_tensor, None
    else:  # full — the activation is split too, so the two w_hi terms take hi and lo separately
        x_hi, x_lo = _split_operand(input_tensor)
        terms = ((x_lo, weight_tensor, None),)

    out = ttnn.experimental.conv3d(
        input_tensor=x_hi, weight_tensor=weight_tensor, bias_tensor=bias_tensor, **conv_kwargs
    )
    for extra_input, extra_weight, extra_bias in ((x_hi, weight_lo_tensor, None),) + (terms or ()):
        out = ttnn.add(
            out,
            ttnn.experimental.conv3d(
                input_tensor=extra_input, weight_tensor=extra_weight, bias_tensor=extra_bias, **conv_kwargs
            ),
        )
    return out


def channel_axis(parallel_config) -> int | None:
    """Mesh axis the channel (C) dim is tensor-parallel over, or None."""
    if isinstance(parallel_config, AudioTCParallelConfig) and parallel_config.channel_parallel.factor > 1:
        return parallel_config.channel_parallel.mesh_axis
    return None


def channel_factor(parallel_config) -> int:
    """Channel-TP factor (per-chip C-shard count = C / factor), or 1 if no C-TP."""
    if channel_axis(parallel_config) is None:
        return 1
    return parallel_config.channel_parallel.factor


def channel_align_unit(parallel_config) -> int:
    """Channel alignment unit: ``factor * TILE_WIDTH`` under channel-TP, else 32."""
    from ..utils.conv3d import ALIGNMENT

    return ALIGNMENT * channel_factor(parallel_config)


def partition_channel(x: ttnn.Tensor, parallel_config, *, dim: int) -> ttnn.Tensor:
    """Local slice of a channel-replicated tensor to its per-chip C-shard (no-op without C-TP)."""
    axis = channel_axis(parallel_config)
    if axis is None:
        return x
    return ttnn.mesh_partition(x, dim=dim, cluster_axis=axis)


def all_gather_channel(ccl_manager, x: ttnn.Tensor, parallel_config, *, dim: int) -> ttnn.Tensor:
    """All-gather a channel-sharded tensor back to full C on the channel axis."""
    axis = channel_axis(parallel_config)
    if axis is None:
        return x
    return ccl_manager.all_gather_persistent_buffer(x, dim=dim, mesh_axis=axis)


def gather_channel_to_full(ccl_manager, x_BTC: ttnn.Tensor, parallel_config) -> ttnn.Tensor:
    """Gather a ``(B, T, C)`` C-shard back to full C_in (no-op without channel-TP)."""
    if channel_axis(parallel_config) is None:
        return x_BTC
    x_BTC = ttnn.to_layout(x_BTC, ttnn.TILE_LAYOUT)
    x_BTC = all_gather_channel(ccl_manager, x_BTC, parallel_config, dim=2)
    return ttnn.to_layout(x_BTC, ttnn.ROW_MAJOR_LAYOUT)


def _pick_c_out_block_shard(*, full: int, shard: int, dst_size: int = 4, tile: int = 32) -> int:
    """Largest C_out_block ≤ full that divides shard and satisfies conv3d matmul_N_t rules."""
    max_nt = min(full, shard) // tile
    for nt in range(max_nt, 0, -1):
        block = nt * tile
        if shard % block != 0:
            continue
        if nt <= dst_size or nt % dst_size == 0:
            return block
    return tile  # 1 tile is always legal


def prepare_conv3d_weight_state(
    state: dict,
    w_5d: torch.Tensor,
    *,
    conv_config,
    mesh_device: ttnn.MeshDevice,
    dtype: ttnn.DataType,
    unpadded_out: int,
    out_channels: int,
    unpadded_in: int | None = None,
    in_channels: int | None = None,
    split: bool = False,
) -> None:
    """Zero-pad the 5D conv weight/bias to aligned size, prepare, and write to ``state``.

    With ``split``, also writes ``state["weight_lo"]``: the weight is decomposed into ``hi = bf16(w)`` and
    the exact residual ``lo = w - hi``, each prepared separately, so `conv3d_maybe_split` can recover the
    mantissa bits the ~11-bit multiplier drops. The padding above is applied once, before the split, so
    the bias is never padded twice.
    """
    if out_channels != unpadded_out:
        pad_co = out_channels - unpadded_out
        w_5d = torch.nn.functional.pad(w_5d, (0, 0, 0, 0, 0, 0, 0, 0, 0, pad_co))
        if "bias" in state:
            state["bias"] = torch.nn.functional.pad(state["bias"], (0, pad_co))
    if in_channels is not None and in_channels != unpadded_in:
        w_5d = torch.nn.functional.pad(w_5d, (0, 0, 0, 0, 0, 0, 0, in_channels - unpadded_in))

    def _prepare(w: torch.Tensor) -> torch.Tensor:
        weight_tt = ttnn.from_torch(w, dtype=dtype, pad_value=0)
        prepared = ttnn.experimental.prepare_conv3d_weights(
            weight_tensor=weight_tt, C_in_block=conv_config.C_in_block, device=mesh_device
        )
        # Replicated, but slicing one coordinate keeps the parent's distribution metadata, which
        # `ttnn.to_torch` refuses on a multi-host mesh; the helper reads a shard this host owns.
        return local_device_to_torch(prepared)

    if split:
        w_hi = w_5d.float().bfloat16().float()
        state["weight"] = _prepare(w_hi)
        state["weight_lo"] = _prepare(w_5d.float() - w_hi)
    else:
        state["weight"] = _prepare(w_5d)


def _t_neighbor_pad(
    x_BTC: ttnn.Tensor,
    *,
    pad_left: int,
    pad_right: int,
    parallel_config: "ParallelFactor | AudioTParallelConfig",
    ccl_manager: CCLManager,
    padding_mode: str = "zeros",
) -> ttnn.Tensor:
    """Halo exchange on the T axis (dim 1 in BTC), single- or two-axis sharded."""
    if pad_left == 0 and pad_right == 0:
        return x_BTC
    if parallel_config is None or parallel_config.factor <= 1:
        B, T, C = x_BTC.shape
        if pad_left > 0:
            zl = ttnn.zeros(
                (B, pad_left, C), dtype=x_BTC.get_dtype(), layout=ttnn.ROW_MAJOR_LAYOUT, device=x_BTC.device()
            )
            x_BTC = ttnn.concat([zl, x_BTC], dim=1)
        if pad_right > 0:
            zr = ttnn.zeros(
                (B, pad_right, C), dtype=x_BTC.get_dtype(), layout=ttnn.ROW_MAJOR_LAYOUT, device=x_BTC.device()
            )
            x_BTC = ttnn.concat([x_BTC, zr], dim=1)
        return x_BTC

    outer_dims = x_BTC.shape[0]
    num_links = max(1, min(outer_dims, ccl_manager.num_links))

    if isinstance(parallel_config, AudioTParallelConfig):
        # Two-axis halo: one call per mesh axis (distinct pad dims required).
        sem0 = ccl_manager.get_np_ping_pong_semaphore(parallel_config.axis0.mesh_axis)
        x_BTC = ccl_manager.neighbor_pad_persistent_buffer(
            x_BTC,
            dims=[1],
            pad_left=[pad_left],
            pad_right=[pad_right],
            padding_mode=padding_mode,
            axes=[parallel_config.axis0.mesh_axis],
            neighbor_sems=[sem0],
            num_links=[num_links],
        )
        sem1 = ccl_manager.get_np_ping_pong_semaphore(parallel_config.axis1.mesh_axis)
        return ccl_manager.neighbor_pad_persistent_buffer(
            x_BTC,
            dims=[1],
            pad_left=[pad_left],
            pad_right=[pad_right],
            padding_mode=padding_mode,
            axes=[parallel_config.axis1.mesh_axis],
            neighbor_sems=[sem1],
            num_links=[num_links],
        )

    sem = ccl_manager.get_np_ping_pong_semaphore(parallel_config.mesh_axis)
    return ccl_manager.neighbor_pad_persistent_buffer(
        x_BTC,
        dims=[1],
        pad_left=[pad_left],
        pad_right=[pad_right],
        padding_mode=padding_mode,
        axes=[parallel_config.mesh_axis],
        neighbor_sems=[sem],
        num_links=[num_links],
    )


_TAP_WARNED: set = set()


def _warn_once(key, message: str) -> None:
    """One warning per distinct key for the process -- the tap filter runs hundreds of times per decode."""
    if key not in _TAP_WARNED:
        _TAP_WARNED.add(key)
        logger.warning(message)


def _tap_weight(taps, channels: int, dtype, mesh_device):
    K = len(taps)
    wt = torch.tensor(taps, dtype=torch.float32).reshape(1, 1, K).expand(channels, 1, K).contiguous()
    return ttnn.from_torch(
        wt,
        layout=ttnn.ROW_MAJOR_LAYOUT,
        dtype=dtype,
        mesh_mapper=ttnn.ReplicateTensorToMesh(mesh_device),
    )


def depthwise_tap_filter(x_BTC, taps, stride, *, mesh_device, dtype, cache):
    """Valid depthwise filter (same K taps per channel) on padded ``(B, T_pad, C)`` ROW_MAJOR.

    Returns ``(B, T_out, C)``, ``T_out = (T_pad - K) / stride + 1``. Runs ``ttnn.conv1d`` (groups=C) with the
    formulation and slice count tabled in ``utils/tap_filter_configs.py``, falling back through the trial chain
    (full C, chunked C, then the bit-equal MAC form) when a row is missing or stale.
    """
    B, T_pad, C = int(x_BTC.shape[0]), int(x_BTC.shape[1]), int(x_BTC.shape[2])
    K = len(taps)
    T_out = (T_pad - K) // stride + 1
    if "cc" not in cache:
        cache["cc"] = ttnn.init_device_compute_kernel_config(
            mesh_device.arch(), math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True, packer_l1_acc=True
        )
    conv_config = ttnn.Conv1dConfig(weights_dtype=dtype, shard_layout=ttnn.TensorMemoryLayout.HEIGHT_SHARDED)
    device_key = tap_device_key(mesh_device)
    shape = f"B={B}, T_pad={T_pad}, C={C}, K={K}, stride={stride}"

    def run(formulation, slice_config):
        if formulation == "mac":
            return _depthwise_tap_mac(x_BTC, taps, stride, T_out=T_out, dtype=dtype)
        channels = C if formulation == "direct" else formulation
        # conv1d prepares the weight for the parallelization it runs with ((B, T_pad) and the slice config); a
        # weight prepared for another one decodes to garbage without raising.
        wkey = ("w", channels, stride, K, tuple(taps), B, T_pad, slice_signature(slice_config))
        weight = cache.get(wkey)
        prepared = weight is not None
        if weight is None:
            weight = _tap_weight(taps, channels, dtype, mesh_device)
        common = dict(
            B=B,
            T_pad=T_pad,
            T_out=T_out,
            K=K,
            stride=stride,
            mesh_device=mesh_device,
            dtype=dtype,
            conv_config=conv_config,
            compute_config=cache["cc"],
            slice_config=slice_config,
            cache=cache,
            wkey=wkey,
        )
        if formulation == "direct":
            return _depthwise_tap_conv1d(x_BTC, weight, C=C, prepared=prepared, **common)
        return _depthwise_tap_conv1d_chunked(x_BTC, weight, C=C, chunk=channels, prepared=prepared, **common)

    attempts: list[tuple] = []  # (formulation, slice config or None for auto, source)
    seen: set = set()

    def add(formulation, slice_config, source):
        key = (formulation, slice_signature(slice_config))
        if key not in seen:
            seen.add(key)
            attempts.append((formulation, slice_config, source))

    shape_key = ("tap_path", B, T_pad, C, K, stride)
    known = cache.get(shape_key)
    if known is not None:
        add(known[0], known[1], "cached")
    else:
        tabled = tap_formulation(device_key, C, K, stride)
        if tabled in applicable_formulations(C):
            channels = C if tabled == "direct" else tabled
            explicit = tap_slice_config(device_key, channels, K, stride, T_out)
            if explicit is not None:
                add(tabled, explicit, "table")
            add(tabled, None, "table")
        else:
            if tabled is not None:
                _warn_once(
                    ("tap_bad_row", device_key, C, K, stride),
                    f"tap filter: table row {format_formulation_row(C, K, stride, tabled)} for {device_key} does not "
                    f"apply to C={C}; ignoring it",
                )
            _warn_once(
                ("tap_miss", device_key, C, K, stride),
                f"tap filter: no table row for {device_key} at C={C}, K={K}, stride={stride}; probing the "
                f"formulations ({shape}). Sweep this shape with tools/sweep_tap_filter_configs.py and add the row "
                f"to utils/tap_filter_configs.py",
            )
    for formulation in applicable_formulations(C) + ["mac"]:
        add(formulation, None, "trial")

    last_exc = None
    for formulation, slice_config, source in attempts:
        try:
            out = run(formulation, slice_config)
        except RuntimeError as exc:
            last_exc = exc
            if source in ("table", "cached"):
                _warn_once(
                    ("tap_stale", device_key, C, K, stride, formulation, slice_signature(slice_config)),
                    f"tap filter: {source} plan {formulation!r} / {slice_signature(slice_config) or 'auto slicing'} "
                    f"failed at {shape}: {str(exc).splitlines()[0][:160]}; falling back",
                )
            continue
        cache[shape_key] = (formulation, slice_config)
        logger.debug(
            f"tap filter plan: {shape} -> {formulation!r} / {slice_signature(slice_config) or 'auto'} ({source})"
        )
        if source == "trial":
            if formulation == "mac":
                _warn_once(
                    ("tap_mac", device_key, B, T_pad, C, K, stride),
                    f"tap filter: no conv1d formulation fits at {shape}; MAC fallback"
                    + (f" ({str(last_exc).splitlines()[0][:120]})" if last_exc is not None else ""),
                )
            else:
                _warn_once(
                    ("tap_probed", device_key, C, K, stride, formulation),
                    f"tap filter: probing found {formulation!r} at {shape}; add "
                    f"{format_formulation_row(C, K, stride, formulation)} to utils/tap_filter_configs.py for "
                    f"{device_key}",
                )
        return out
    raise last_exc


def _depthwise_tap_conv1d_chunked(
    x_BTC,
    weight,
    *,
    B,
    C,
    T_pad,
    T_out,
    K,
    stride,
    chunk,
    mesh_device,
    dtype,
    conv_config,
    compute_config,
    slice_config,
    cache,
    wkey,
    prepared,
):
    """``_depthwise_tap_conv1d`` over independent ``chunk``-wide channel slices, concatenated back.

    Slicing C is exact for a depthwise filter -- no partial sums to reconcile. The reassembly is a
    last-dim concat, lossy in fp32 unless the row is a multiple of the 64B buffer alignment; every chunk
    here is 32 fp32 elements = 128B, and the assert keeps a future chunk size from silently truncating
    the mantissa to TF32.
    """
    assert (chunk * 4) % 64 == 0, f"C-chunk {chunk} would make ttnn.concat(dim=-1) lossy in fp32"
    parts = []
    for start in range(0, C, chunk):
        piece = ttnn.slice(x_BTC, [0, 0, start], [B, T_pad, start + chunk])
        parts.append(
            _depthwise_tap_conv1d(
                piece,
                weight,
                B=B,
                C=chunk,
                T_pad=T_pad,
                T_out=T_out,
                K=K,
                stride=stride,
                mesh_device=mesh_device,
                dtype=dtype,
                conv_config=conv_config,
                compute_config=compute_config,
                slice_config=slice_config,
                cache=cache,
                wkey=wkey,
                prepared=prepared or start > 0,
            )
        )
        ttnn.deallocate(piece)
    return ttnn.concat(parts, dim=2)


def _depthwise_tap_mac(x_BTC, taps, stride, *, T_out: int, dtype):
    """Depthwise K-tap filter as ``sum_k tap_k * x[k : k + stride*T_out : stride]``.

    No convolution op, so no shard-layout or slicing constraints. Used when ``ttnn.conv1d``
    cannot find a valid configuration for the shape. The result is cast to ``dtype`` so both
    forms return the caller's requested conv dtype.
    """
    accumulator = None
    for offset, tap in enumerate(taps):
        if tap == 0.0:
            continue
        window = ttnn.slice(
            x_BTC,
            [0, offset, 0],
            [x_BTC.shape[0], offset + stride * (T_out - 1) + 1, x_BTC.shape[2]],
            [1, stride, 1],
        )
        scaled = ttnn.multiply(window, float(tap))
        accumulator = scaled if accumulator is None else ttnn.add(accumulator, scaled)
    if accumulator is not None and accumulator.get_dtype() != dtype:
        accumulator = ttnn.typecast(accumulator, dtype)
    return accumulator


def _depthwise_tap_conv1d(
    x_BTC,
    weight,
    *,
    B,
    C,
    T_pad,
    T_out,
    K,
    stride,
    mesh_device,
    dtype,
    conv_config,
    compute_config,
    slice_config,
    cache,
    wkey,
    prepared,
):
    """The HEIGHT_SHARDED ``ttnn.conv1d`` fast path; ``slice_config`` None leaves the DRAM slicing to the op."""
    out, _, (weight, _bias) = ttnn.conv1d(
        input_tensor=ttnn.reshape(x_BTC, (B, T_pad, 1, C)),
        weight_tensor=weight,
        device=mesh_device,
        in_channels=C,
        out_channels=C,
        batch_size=B,
        input_length=T_pad,
        kernel_size=K,
        stride=stride,
        padding=0,
        dilation=1,
        groups=C,
        dtype=dtype,
        conv_config=conv_config,
        compute_config=compute_config,
        slice_config=slice_config,
        return_output_dim=True,
        return_weights_and_bias=True,
    )
    if not prepared:
        cache[wkey] = weight
    # conv1d emits HEIGHT_SHARDED TILE; untilize straight to interleaved DRAM ROW_MAJOR — the
    # memory_config folds the reshard into the untilize (one op instead of to_memory_config + to_layout).
    out = ttnn.to_layout(out, ttnn.ROW_MAJOR_LAYOUT, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    return ttnn.reshape(out, (B, T_out, C))


def _all_gather_t(ccl_manager, x: "ttnn.Tensor", parallel_config) -> "ttnn.Tensor":
    """All-gather the T-sharded tensor to full T on every chip."""
    if isinstance(parallel_config, AudioTParallelConfig):
        x = ccl_manager.all_gather_persistent_buffer(x, dim=1, mesh_axis=parallel_config.axis1.mesh_axis)
        x = ccl_manager.all_gather_persistent_buffer(x, dim=1, mesh_axis=parallel_config.axis0.mesh_axis)
    else:
        x = ccl_manager.all_gather_persistent_buffer(x, dim=1, mesh_axis=parallel_config.mesh_axis)
    return x


# T-shard alignment is `factor`, not `32 * factor`: the coarser alignment only existed because
# TILE-layout mesh_partition needs a tile-aligned offset, but partitioning is done in ROW_MAJOR
# (fine at any offset) and every caller untilized right after. The tighter pad (49 -> 1 rows at
# T=207, factor 8) also lets `_set_tpad_tail` fix just the pad suffix in place (`_set_tpad_tail_local`).


def _partition_t(x: "ttnn.Tensor", parallel_config) -> "ttnn.Tensor":
    """Partition T across the mesh (inverse of _all_gather_t)."""
    if isinstance(parallel_config, AudioTParallelConfig):
        x = ttnn.mesh_partition(x, dim=1, cluster_axis=parallel_config.axis0.mesh_axis)
        x = ttnn.mesh_partition(x, dim=1, cluster_axis=parallel_config.axis1.mesh_axis)
    else:
        x = ttnn.mesh_partition(x, dim=1, cluster_axis=parallel_config.mesh_axis)
    return x


def _make_kaiser_sinc_kernel_1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
    """Return a shape-``(kernel_size,)`` kaiser-windowed sinc filter."""
    even = kernel_size % 2 == 0
    half_size = kernel_size // 2
    delta_f = 4 * half_width
    amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
    if amplitude > 50.0:
        beta = 0.1102 * (amplitude - 8.7)
    elif amplitude >= 21.0:
        beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
    else:
        beta = 0.0
    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
    if even:
        time = torch.arange(-half_size, half_size) + 0.5
    else:
        time = torch.arange(kernel_size) - half_size
    if cutoff == 0:
        filter_ = torch.zeros_like(time)
    else:
        filter_ = (
            2
            * cutoff
            * window
            * torch.where(
                time == 0,
                torch.tensor(1.0, dtype=time.dtype),
                torch.sin(math.pi * 2 * cutoff * time) / (math.pi * 2 * cutoff * time),
            )
        )
        filter_ = filter_ / filter_.sum()
    return filter_.float().reshape(kernel_size)


def _replicate_pad_t(x_BTC: ttnn.Tensor, pad_left: int, pad_right: int, mesh_device: ttnn.MeshDevice) -> ttnn.Tensor:
    """Replicate-pad along the T axis for a ``(B, T, C)`` ROW_MAJOR tensor."""
    if pad_left == 0 and pad_right == 0:
        return x_BTC
    B, T, C = x_BTC.shape
    pieces = []
    if pad_left > 0:
        first = ttnn.slice(x_BTC, [0, 0, 0], [B, 1, C])
        pieces.extend([first] * pad_left)
    pieces.append(x_BTC)
    if pad_right > 0:
        last = ttnn.slice(x_BTC, [0, T - 1, 0], [B, T, C])
        pieces.extend([last] * pad_right)
    return ttnn.concat(pieces, dim=1)


def _tpad_mask(mesh_device, parallel_config, dtype, global_T, tpad_image, cache):
    """Cached T-sharded validity mask ``M`` (1.0 real, 0.0 on trailing ``tpad_image`` rows) and its complement ``inv``."""
    key = (global_T, tpad_image, dtype)
    cached = cache.get(key)
    if cached is None:
        m = torch.ones(1, global_T, 1, dtype=torch.float32)
        m[:, global_T - tpad_image :, :] = 0.0
        pair = []
        for t in (m, 1.0 - m):
            # The mask ends up ROW_MAJOR anyway, so build it that way and skip the tile-aligned
            # partition constraint entirely.
            mt = ttnn.from_torch(t, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=dtype)
            pair.append(_partition_t(mt, parallel_config))
        cached = tuple(pair)
        cache[key] = cached
    return cached


def _tpad_mask_suffix(mesh_device, parallel_config, dtype, global_T, tpad_image, local_T, cache):
    """The last ``tpad_image`` local rows of the validity mask and its complement.

    Only the final shard(s) carry pad rows, so away from them the suffix mask is all ones and
    multiplying by it is a no-op -- which is what lets one mesh-uniform local range stand in for a
    pad region that sits at a different global offset on every shard.
    """
    key = ("suffix", global_T, tpad_image, local_T, dtype)
    cached = cache.get(key)
    if cached is None:
        m = torch.ones(1, global_T, 1, dtype=torch.float32)
        m[:, global_T - tpad_image :, :] = 0.0
        # Reshape to per-shard rows, keep each shard's trailing `tpad_image` rows, and restack so the
        # mesh partition below hands every device its own suffix.
        per_shard = m.reshape(1, global_T // local_T, local_T, 1)[:, :, local_T - tpad_image :, :]
        suffix = per_shard.reshape(1, -1, 1)
        pair = []
        for t in (suffix, 1.0 - suffix):
            mt = ttnn.from_torch(t, device=mesh_device, layout=ttnn.ROW_MAJOR_LAYOUT, dtype=dtype)
            pair.append(_partition_t(mt, parallel_config))
        cached = tuple(pair)
        cache[key] = cached
    return cached


def _set_tpad_tail(x_BTC, tpad_image, *, mode, mesh_device, parallel_config, cache, ccl_manager=None):
    """Set the trailing ``tpad_image`` tail rows: ``mode="zeros"`` zeros them, ``mode="replicate"`` fills the last real row.

    Uses a cached validity mask (body rows multiply by 1.0, staying bit-identical). CCL-free except when the
    pad image spans several shards in replicate mode, where one 32-row gather fetches the holder shard's row.
    """
    if tpad_image <= 0 or parallel_config is None or getattr(parallel_config, "factor", 0) <= 1:
        return x_BTC
    local_T = x_BTC.shape[1]
    if tpad_image < local_T:
        # The pad fits inside one shard's tail, so fix just those rows in place instead of
        # multiplying every body row by 1.0.
        return _set_tpad_tail_local(
            x_BTC, tpad_image, mode=mode, mesh_device=mesh_device, parallel_config=parallel_config, cache=cache
        )
    M, inv = _tpad_mask(
        mesh_device, parallel_config, x_BTC.get_dtype(), local_T * parallel_config.factor, tpad_image, cache
    )
    xm = ttnn.multiply(x_BTC, M)
    if mode == "zeros":
        return xm
    if mode != "replicate":
        raise ValueError(f"unknown mode {mode!r}")
    # The real-last row's LOCAL offset is uniform across shards, but only the holder shard's row at that
    # offset is real -- on a fully-padded shard it is a pad row. Gather a tile-aligned 32-row window holding
    # that offset along the shard axis and take the holder's row, so every shard replicates the true last
    # real row. (A 1-row gather, tile-padded or row-major, hangs on a 32-chip axis; a full-tile window works.)
    global_T = local_T * parallel_config.factor
    assert tpad_image < global_T, f"pad image ({tpad_image}) leaves no real rows (global T {global_T})"
    assert ccl_manager is not None, "multi-shard replicate tail needs ccl_manager to fetch the holder shard's row"
    tile = ttnn.TILE_SIZE
    assert local_T >= tile, f"multi-shard replicate tail needs >= {tile} rows per shard, got {local_T}"
    last_global = global_T - tpad_image - 1
    idx, holder = last_global % local_T, last_global // local_T
    start = min(idx, local_T - tile)
    B, C = x_BTC.shape[0], x_BTC.shape[2]
    window = ttnn.slice(x_BTC, [0, start, 0], [B, start + tile, C])
    window_t = ttnn.to_layout(window, ttnn.TILE_LAYOUT)
    ttnn.deallocate(window)
    # `_all_gather_t` hands back the CCL manager's persistent buffer -- convert, never deallocate it.
    gathered = ttnn.to_layout(_all_gather_t(ccl_manager, window_t, parallel_config), ttnn.ROW_MAJOR_LAYOUT)
    ttnn.deallocate(window_t)
    row = holder * tile + (idx - start)
    last = ttnn.slice(gathered, [0, row, 0], [B, row + 1, C])
    ttnn.deallocate(gathered)
    fill = ttnn.multiply(last, inv)
    ttnn.deallocate(last)
    out = ttnn.add(xm, fill)
    ttnn.deallocate(xm)
    ttnn.deallocate(fill)
    return out


def _set_tpad_tail_local(x_BTC, tpad_image, *, mode, mesh_device, parallel_config, cache):
    """``_set_tpad_tail`` restricted to the pad suffix, written back in place.

    Slices the suffix, fixes it, and writes it back with `ttnn.experimental.slice_write`
    (bit-identical to the full-tensor multiply, cheaper in ROW_MAJOR). Safe in place because every
    op with a receptive field runs its own tail-set first (no stale reads) and the pad is cropped
    from the final waveform. The caller's tensor is mutated and returned, so its
    `if xs is not x: deallocate(xs)` correctly skips the free.
    """
    B, local_T, C = x_BTC.shape
    global_T = local_T * parallel_config.factor
    assert tpad_image < global_T, f"pad image ({tpad_image}) leaves no real rows (global T {global_T})"
    if mode not in ("zeros", "replicate"):
        raise ValueError(f"unknown mode {mode!r}")
    begin = local_T - tpad_image
    M_suf, inv_suf = _tpad_mask_suffix(
        mesh_device, parallel_config, x_BTC.get_dtype(), global_T, tpad_image, local_T, cache
    )

    suffix = ttnn.slice(x_BTC, [0, begin, 0], [B, local_T, C])
    fixed = ttnn.multiply(suffix, M_suf)
    ttnn.deallocate(suffix)
    if mode == "replicate":
        # The real-last row sits just before the pad on the shard that holds it; on every other shard
        # inv_suf is zero, so whichever row this picks up there contributes nothing.
        idx = (global_T - tpad_image - 1) % local_T
        last = ttnn.slice(x_BTC, [0, idx, 0], [B, idx + 1, C])
        fill = ttnn.multiply(last, inv_suf)
        ttnn.deallocate(last)
        combined = ttnn.add(fixed, fill)
        ttnn.deallocate(fixed)
        ttnn.deallocate(fill)
        fixed = combined
    ttnn.experimental.slice_write(fixed, x_BTC, [0, begin, 0], [B, local_T, C], [1, 1, 1])
    ttnn.deallocate(fixed)
    return x_BTC


def _persistent_zeros(shape, *, dtype, layout, mesh_device: ttnn.MeshDevice) -> ttnn.Tensor:
    """Cached read-only zeros constant for (shape, dtype, layout); avoids per-call host→device writes so traces capture."""
    cache = _ZEROS_CACHE.setdefault(id(mesh_device), {})
    key = (tuple(shape), dtype, layout)
    z = cache.get(key)
    if z is None:
        z = ttnn.zeros(shape, dtype=dtype, layout=layout, device=mesh_device)
        cache[key] = z
    return z


def _zero_pad_t(x_BTC: ttnn.Tensor, pad_left: int, pad_right: int, mesh_device: ttnn.MeshDevice) -> ttnn.Tensor:
    """Zero-pad along the T axis."""
    if pad_left == 0 and pad_right == 0:
        return x_BTC
    B, T, C = x_BTC.shape
    pieces = []
    dtype = x_BTC.get_dtype()
    if pad_left > 0:
        zeros = _persistent_zeros((B, pad_left, C), dtype=dtype, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=mesh_device)
        pieces.append(zeros)
    pieces.append(x_BTC)
    if pad_right > 0:
        zeros = _persistent_zeros((B, pad_right, C), dtype=dtype, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=mesh_device)
        pieces.append(zeros)
    return ttnn.concat(pieces, dim=1)


def _pad_channels_to_aligned(x_BTC: ttnn.Tensor, mesh_device: ttnn.MeshDevice, channel_align: int = 32) -> ttnn.Tensor:
    """Pad C up to ``aligned_channels(C, channel_align)`` with zeros. No-op if aligned."""
    B, T, C = x_BTC.shape
    aligned = aligned_channels(C, channel_align)
    if aligned == C:
        return x_BTC
    pad_c = aligned - C
    dtype = x_BTC.get_dtype()
    zeros = _persistent_zeros((B, T, pad_c), dtype=dtype, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=mesh_device)
    return ttnn.concat([x_BTC, zeros], dim=2)


def _zero_stuff_t(x_BTC: ttnn.Tensor, *, stride: int, mesh_device: ttnn.MeshDevice) -> ttnn.Tensor:
    """Insert ``stride-1`` zeros between T samples (output ``T*s - (s-1)``) via concat+reshape."""
    if stride == 1:
        return x_BTC
    B, T, C = x_BTC.shape
    dtype = x_BTC.get_dtype()
    x_btoc = ttnn.reshape(x_BTC, (B, T, 1, C))
    zero_block = _persistent_zeros(
        (B, T, stride - 1, C), dtype=dtype, layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=mesh_device
    )
    stacked = ttnn.concat([x_btoc, zero_block], dim=2)
    interleaved = ttnn.reshape(stacked, (B, T * stride, C))
    out_len = T * stride - (stride - 1)
    return ttnn.slice(interleaved, [0, 0, 0], [B, out_len, C])


class Conv2dViaConv3d(Module):
    """2D conv via conv3d (kernel ``(1, kh, kw)``) on ``(B, 1, H, W, C)``, single-device.

    Padding: ``"zeros"`` (symmetric internal) or ``"causal_height"``/``"causal_width"`` (external front pad).
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        kernel_size: int | Sequence[int],
        stride: int | Sequence[int] = 1,
        padding_mode: str = "zeros",
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.bfloat16,
        split_mode: str = "off",
    ) -> None:
        super().__init__()

        if padding_mode not in ("zeros", "causal_height", "causal_width"):
            raise ValueError(f"padding_mode must be zeros/causal_height/causal_width, got {padding_mode!r}")
        if split_mode not in CONV_SPLIT_MODES:
            raise ValueError(f"split_mode must be one of {CONV_SPLIT_MODES}, got {split_mode!r}")

        self.unpadded_in_channels = in_channels
        self.unpadded_out_channels = out_channels
        self.in_channels = aligned_channels(in_channels)
        self.out_channels = max(32, out_channels)
        if self.out_channels != self.unpadded_out_channels:
            logger.warning(f"Padding out_channels from {self.unpadded_out_channels} to {self.out_channels}")

        kh, kw = _ntuple(kernel_size, 2)
        sh, sw = _ntuple(stride, 2)
        self.kernel_size = (1, kh, kw)
        self.stride = (1, sh, sw)
        self.padding_mode = padding_mode
        self.pad_h = kh - 1
        self.pad_w = kw - 1
        self.mesh_device = mesh_device
        self.dtype = dtype

        # Causal modes pad externally (internal pad=0 on that axis); zeros pads symmetrically.
        if padding_mode == "zeros":
            self.internal_padding = (0, kh // 2, kw // 2)
        elif padding_mode == "causal_height":
            self.internal_padding = (0, 0, kw // 2)
        else:  # causal_width
            self.internal_padding = (0, kh // 2, 0)

        self.conv_config = get_conv3d_config(
            self.in_channels,
            self.out_channels,
            self.kernel_size,
            dtype,
            grid_size=self.mesh_device.compute_with_storage_grid_size(),
            h_factor=1,
            w_factor=1,
        )

        from models.common.utility_functions import is_blackhole

        self.compute_kernel_config = ttnn.init_device_compute_kernel_config(
            self.mesh_device.arch(),
            math_fidelity=(
                ttnn.MathFidelity.HiFi4 if (is_blackhole() or dtype == ttnn.float32) else ttnn.MathFidelity.HiFi2
            ),
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=False,
        )

        # An explicit constructor argument: MiniMax-H3 opts in, LTX keeps the "off" default. See
        # `conv3d_maybe_split`; splitting only helps an fp32 datapath.
        self.split_mode = split_mode if dtype == ttnn.float32 else "off"

        d = self.kernel_size[0] * self.kernel_size[1] * self.kernel_size[2] * self.in_channels
        self.weight = Parameter(total_shape=[d, self.out_channels], device=mesh_device, pad_value=0, dtype=dtype)
        self.weight_lo = (
            Parameter(total_shape=[d, self.out_channels], device=mesh_device, pad_value=0, dtype=dtype)
            if self.split_mode != "off"
            else None
        )
        self.bias = Parameter(total_shape=[1, self.out_channels], device=mesh_device, pad_value=0, dtype=dtype)

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        """Map ``Conv2d.weight (Cout, Cin, kh, kw)`` → conv3d-prepared ``(d, Cout)``."""
        if "conv.weight" in state:
            state["weight"] = state.pop("conv.weight")
        if "conv.bias" in state:
            state["bias"] = state.pop("conv.bias")

        if "weight" in state:
            w = state["weight"]
            assert w.dim() == 4, f"expected 4D Conv2d weight, got {tuple(w.shape)}"
            prepare_conv3d_weight_state(
                state,
                w.unsqueeze(2).contiguous(),  # → (Cout, Cin, 1, kh, kw)
                conv_config=self.conv_config,
                mesh_device=self.mesh_device,
                dtype=self.dtype,
                unpadded_out=self.unpadded_out_channels,
                out_channels=self.out_channels,
                split=self.split_mode != "off",
            )
        if "bias" in state:
            state["bias"] = state["bias"].reshape(1, -1)

    def forward(self, x_BHWC: ttnn.Tensor) -> ttnn.Tensor:
        """``x_BHWC``: ``(B, H, W, C)`` ROW_MAJOR. Internally pads to 5D and calls conv3d."""
        assert x_BHWC.layout == ttnn.ROW_MAJOR_LAYOUT, f"expected ROW_MAJOR, got {x_BHWC.layout}"
        B, H, W, C = x_BHWC.shape

        if self.padding_mode == "causal_height" and self.pad_h > 0:
            B_, H_, W_, C_ = x_BHWC.shape
            pad_tensor_shape = (B_, self.pad_h, W_, C_)
            zero_pad = _persistent_zeros(
                pad_tensor_shape, dtype=x_BHWC.get_dtype(), layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=self.mesh_device
            )
            x_BHWC = ttnn.concat([zero_pad, x_BHWC], dim=1)
        elif self.padding_mode == "causal_width" and self.pad_w > 0:
            B_, H_, W_, C_ = x_BHWC.shape
            pad_tensor_shape = (B_, H_, self.pad_w, C_)
            zero_pad = _persistent_zeros(
                pad_tensor_shape, dtype=x_BHWC.get_dtype(), layout=ttnn.ROW_MAJOR_LAYOUT, mesh_device=self.mesh_device
            )
            x_BHWC = ttnn.concat([zero_pad, x_BHWC], dim=2)

        x_5d = ttnn.reshape(x_BHWC, (x_BHWC.shape[0], 1, x_BHWC.shape[1], x_BHWC.shape[2], x_BHWC.shape[3]))

        out_5d = conv3d_maybe_split(
            split_mode=self.split_mode,
            input_tensor=x_5d,
            weight_tensor=self.weight.data,
            weight_lo_tensor=self.weight_lo.data if self.weight_lo is not None else None,
            bias_tensor=self.bias.data,
            config=self.conv_config,
            output_channels=self.out_channels,
            kernel_size=self.kernel_size,
            stride=self.stride,
            padding=self.internal_padding,
            padding_mode="zeros",
            dtype=self.dtype,
            compute_kernel_config=self.compute_kernel_config,
        )

        out = ttnn.reshape(out_5d, (out_5d.shape[0], out_5d.shape[2], out_5d.shape[3], out_5d.shape[4]))
        return out


class Conv1dViaConv3d(Module):
    """1D conv via conv3d (kernel ``(k, 1, 1)``) on ``(B, T, 1, 1, C)``.

    Padding: ``"zeros"`` (symmetric "same") or ``"causal"`` (front pad on T); T-sharding turns the pad into a halo exchange.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        kernel_size: int,
        stride: int = 1,
        dilation: int = 1,
        padding: int | None = None,
        padding_mode: str = "zeros",
        bias: bool = True,
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.float32,
        parallel_config: ParallelFactor | None = None,
        ccl_manager: CCLManager | None = None,
        channel_shard_output: bool = True,
        split_mode: str = "off",
    ) -> None:
        super().__init__()

        if padding_mode not in ("zeros", "causal"):
            raise ValueError(f"padding_mode must be zeros/causal, got {padding_mode!r}")
        if split_mode not in CONV_SPLIT_MODES:
            raise ValueError(f"split_mode must be one of {CONV_SPLIT_MODES}, got {split_mode!r}")
        sharded = parallel_config is not None and parallel_config.factor > 1
        if sharded:
            assert ccl_manager is not None, "T-sharding requires ccl_manager"
        self.channel_shard_output = channel_shard_output
        if channel_axis(parallel_config) is not None:
            assert ccl_manager is not None, "channel-TP requires ccl_manager"

        self.channel_align = channel_align_unit(parallel_config)
        self.unpadded_in_channels = in_channels
        self.unpadded_out_channels = out_channels
        self.in_channels = aligned_channels(in_channels, self.channel_align)
        self.out_channels = aligned_channels(max(32, out_channels), self.channel_align)
        if self.out_channels != self.unpadded_out_channels:
            logger.warning(f"Padding out_channels from {self.unpadded_out_channels} to {self.out_channels}")

        self.kernel_size = (kernel_size, 1, 1)
        self.stride = (stride, 1, 1)
        self.dilation = dilation
        self.padding_mode = padding_mode
        self.bias_enabled = bias
        self.mesh_device = mesh_device
        self.dtype = dtype
        self.parallel_config = parallel_config
        self.ccl_manager = ccl_manager

        eff_k = (kernel_size - 1) * dilation + 1
        # ``eff_k // 2`` (the ``padding=None`` default) is torch's "same" padding and is right
        # for every LTX call site. A strided conv need not use it: MiniMax-H3's DAC encoder pairs
        # ``kernel = 2 * stride`` with ``padding = ceil(stride / 2)``, which yields exactly
        # ``L / stride`` where ``eff_k // 2`` would yield ``L / stride + 1``.
        same_pad = eff_k // 2 if padding is None else padding

        if sharded:
            # Zero internal T pad; halo exchange in forward() supplies context.
            self.internal_padding = (0, 0, 0)
            self.external_pad_front = 0
            if padding_mode == "zeros":
                self.halo_pad_left = same_pad
                self.halo_pad_right = same_pad
            else:  # causal
                self.halo_pad_left = eff_k - 1
                self.halo_pad_right = 0
        elif padding_mode == "zeros":
            self.internal_padding = (same_pad, 0, 0)
            self.external_pad_front = 0
            self.halo_pad_left = 0
            self.halo_pad_right = 0
        else:  # causal — pad k-1 at front externally
            self.internal_padding = (0, 0, 0)
            self.external_pad_front = eff_k - 1
            self.halo_pad_left = 0
            self.halo_pad_right = 0

        self.conv_config = get_conv3d_config(
            self.in_channels,
            self.out_channels,
            self.kernel_size,
            dtype,
            grid_size=self.mesh_device.compute_with_storage_grid_size(),
            h_factor=1,
            w_factor=1,
        )

        # Column-parallel C-TP: each chip owns out_channels/factor C_out; C_in stays full (gathered).
        self.out_channels_shard = self.out_channels // channel_factor(parallel_config)
        if channel_factor(parallel_config) > 1:
            self.conv_config_shard = ttnn.Conv3dConfig(
                weights_dtype=self.conv_config.weights_dtype,
                output_layout=ttnn.ROW_MAJOR_LAYOUT,
                T_out_block=self.conv_config.T_out_block,
                W_out_block=self.conv_config.W_out_block,
                H_out_block=self.conv_config.H_out_block,
                C_out_block=_pick_c_out_block_shard(full=self.conv_config.C_out_block, shard=self.out_channels_shard),
                C_in_block=self.conv_config.C_in_block,
                compute_with_storage_grid_size=self.mesh_device.compute_with_storage_grid_size(),
            )

        self.compute_kernel_config = ttnn.init_device_compute_kernel_config(
            self.mesh_device.arch(),
            math_fidelity=ttnn.MathFidelity.HiFi4,
            math_approx_mode=False,
            fp32_dest_acc_en=True,
            packer_l1_acc=True,
        )

        # An explicit constructor argument: MiniMax-H3 opts in, LTX's audio path keeps the fast
        # default. Splitting only helps an fp32 datapath.
        self.split_mode = split_mode if dtype == ttnn.float32 else "off"

        self.same_pad = same_pad
        self.eff_k = eff_k

        self._alloc_weight_bias()

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        """Map torch ``Conv1d.weight (Cout, Cin, k)`` → conv3d-prepared ``(d, Cout)``."""
        if "conv.weight" in state:
            state["weight"] = state.pop("conv.weight")
        if "conv.bias" in state:
            state["bias"] = state.pop("conv.bias")

        if "weight" in state:
            w = state["weight"]
            assert w.dim() == 3, f"expected 3D Conv1d weight, got {tuple(w.shape)}"
            prepare_conv3d_weight_state(
                state,
                w.unsqueeze(-1).unsqueeze(-1).contiguous(),  # → (Cout, Cin, k, 1, 1)
                conv_config=self.conv_config,
                mesh_device=self.mesh_device,
                dtype=self.dtype,
                unpadded_out=self.unpadded_out_channels,
                out_channels=self.out_channels,
                unpadded_in=self.unpadded_in_channels,
                in_channels=self.in_channels,
                split=self.split_mode != "off",
            )
        if "bias" in state and self.bias is not None:
            state["bias"] = state["bias"].reshape(1, -1)

    def _is_col_parallel(self) -> bool:
        return channel_axis(self.parallel_config) is not None and self.channel_shard_output

    def _alloc_weight_bias(self) -> None:
        """Allocate weight/bias; column-parallel shards C_out across the channel axis at load."""
        d = self.kernel_size[0] * self.kernel_size[1] * self.kernel_size[2] * self.in_channels
        mesh_axes = [None, channel_axis(self.parallel_config)] if self._is_col_parallel() else None

        self.weight = Parameter(
            total_shape=[d, self.out_channels],
            device=self.mesh_device,
            pad_value=0,
            dtype=self.dtype,
            mesh_axes=mesh_axes,
        )
        # The weight residual is laid out exactly like the weight, so it shards the same way.
        self.weight_lo = (
            Parameter(
                total_shape=[d, self.out_channels],
                device=self.mesh_device,
                pad_value=0,
                dtype=self.dtype,
                mesh_axes=mesh_axes,
            )
            if self.split_mode != "off"
            else None
        )
        self.bias = (
            Parameter(
                total_shape=[1, self.out_channels],
                device=self.mesh_device,
                pad_value=0,
                dtype=self.dtype,
                mesh_axes=mesh_axes,
            )
            if self.bias_enabled
            else None
        )

    def _conv_args(self):
        """``(weight, weight_lo, bias, conv_config, output_channels)``; column-parallel returns the per-chip C_out shard."""
        bias = self.bias.data if self.bias is not None else None
        weight_lo = self.weight_lo.data if self.weight_lo is not None else None
        if self._is_col_parallel():
            return self.weight.data, weight_lo, bias, self.conv_config_shard, self.out_channels_shard
        return self.weight.data, weight_lo, bias, self.conv_config, self.out_channels

    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
        """``x_BTC``: ``(B, T, C)`` ROW_MAJOR → ``(B, T_out, C_out)`` (T sharded uses a halo exchange)."""
        assert x_BTC.layout == ttnn.ROW_MAJOR_LAYOUT

        # Channel-TP: gather C_in to full, then column-parallel emits this chip's C_out slice.
        x_BTC = gather_channel_to_full(self.ccl_manager, x_BTC, self.parallel_config)
        weight, weight_lo, bias, conv_config, out_channels = self._conv_args()

        if self.parallel_config is not None and self.parallel_config.factor > 1:
            x_BTC = _t_neighbor_pad(
                x_BTC,
                pad_left=self.halo_pad_left,
                pad_right=self.halo_pad_right,
                parallel_config=self.parallel_config,
                ccl_manager=self.ccl_manager,
                padding_mode="zeros",
            )
        elif self.external_pad_front > 0:
            B, T, C = x_BTC.shape
            zero_pad = ttnn.zeros(
                (B, self.external_pad_front, C),
                dtype=x_BTC.get_dtype(),
                layout=ttnn.ROW_MAJOR_LAYOUT,
                device=self.mesh_device,
            )
            x_BTC = ttnn.concat([zero_pad, x_BTC], dim=1)

        x_5d = ttnn.reshape(x_BTC, (x_BTC.shape[0], x_BTC.shape[1], 1, 1, x_BTC.shape[2]))

        out_5d = conv3d_maybe_split(
            split_mode=self.split_mode,
            input_tensor=x_5d,
            weight_tensor=weight,
            weight_lo_tensor=weight_lo,
            bias_tensor=bias,
            config=conv_config,
            output_channels=out_channels,
            kernel_size=self.kernel_size,
            stride=self.stride,
            padding=self.internal_padding,
            dilation=(self.dilation, 1, 1),
            padding_mode="zeros",
            dtype=self.dtype,
            compute_kernel_config=self.compute_kernel_config,
        )
        return ttnn.reshape(out_5d, (out_5d.shape[0], out_5d.shape[1], out_5d.shape[4]))


class _AlignedOutConv1d(Conv1dViaConv3d):
    """Conv1dViaConv3d variant that rounds ``out_channels`` to a 32-multiple, then trims back to the real count."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        kernel_size: int,
        stride: int = 1,
        dilation: int = 1,
        padding: int | None = None,
        padding_mode: str = "zeros",
        bias: bool = True,
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.float32,
        parallel_config: ParallelFactor | None = None,
        ccl_manager: CCLManager | None = None,
        channel_shard_output: bool = True,
        split_mode: str = "off",
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            dilation=dilation,
            padding=padding,
            padding_mode=padding_mode,
            bias=bias,
            mesh_device=mesh_device,
            dtype=dtype,
            parallel_config=parallel_config,
            ccl_manager=ccl_manager,
            channel_shard_output=channel_shard_output,
            split_mode=split_mode,
        )

    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
        # Channel-TP gathers C_in in super().forward, so skip the local pad here.
        if channel_axis(self.parallel_config) is None:
            x_BTC = _pad_channels_to_aligned(x_BTC, self.mesh_device, channel_align=self.channel_align)
        y = super().forward(x_BTC)
        # Column-parallel can't trim per chip; the trim happens once at the output.
        if not self._is_col_parallel() and self.unpadded_out_channels < self.out_channels:
            B, T, C = y.shape
            y = ttnn.slice(y, [0, 0, 0], [B, T, self.unpadded_out_channels])
        return y


class ConvTranspose1dViaConv3d(Module):
    """Substitute for ``torch.nn.ConvTranspose1d`` (``padding=(k-stride)//2``): Conv1d on zero-stuffed input with a flipped/transposed weight."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        kernel_size: int,
        stride: int,
        bias: bool = True,
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.float32,
        parallel_config: ParallelFactor | None = None,
        ccl_manager: CCLManager | None = None,
        split_mode: str = "off",
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.external_pad_each = kernel_size - 1 - (kernel_size - stride) // 2
        self.mesh_device = mesh_device
        self.dtype = dtype
        self.parallel_config = parallel_config
        self.ccl_manager = ccl_manager

        # Inner conv stays UNSHARDED; forward gathers T, runs unsharded, then re-partitions.
        self.conv = _AlignedOutConv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=1,
            padding_mode="causal",
            bias=bias,
            mesh_device=mesh_device,
            dtype=dtype,
            parallel_config=None,
            ccl_manager=None,
            split_mode=split_mode,
        )
        # forward() supplies its own symmetric padding, so the inner conv's causal front pad is disabled.
        self.conv.external_pad_front = 0

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        """Reshape ConvTranspose1d weight ``(in, out, k)`` → Conv1d ``(out, in, k)`` (flipped along k)."""
        if "weight" in state:
            w = state.pop("weight")
            assert w.dim() == 3 and tuple(w.shape) == (self.in_channels, self.out_channels, self.kernel_size), (
                f"expected ConvTranspose1d weight shape ({self.in_channels}, {self.out_channels}, "
                f"{self.kernel_size}), got {tuple(w.shape)}"
            )
            w_flipped = torch.flip(w, dims=[-1])
            w_conv1d = w_flipped.permute(1, 0, 2).contiguous()
            state["conv.weight"] = w_conv1d
        if "bias" in state:
            state["conv.bias"] = state.pop("bias")

    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
        assert x_BTC.layout == ttnn.ROW_MAJOR_LAYOUT
        sharded = self.parallel_config is not None and self.parallel_config.factor > 1

        # Channel-TP: gather C_in here, scatter C_out at the end (inner conv is unsharded).
        ch_axis = channel_axis(self.parallel_config)
        if ch_axis is not None:
            x_BTC = gather_channel_to_full(self.ccl_manager, x_BTC, self.parallel_config)
            # Drop the gathered pad channels so the aligned-32 inner conv sees its real C_in.
            x_BTC = ttnn.slice(x_BTC, [0, 0, 0], [x_BTC.shape[0], x_BTC.shape[1], self.in_channels])

        if sharded:
            x_BTC = ttnn.to_layout(x_BTC, ttnn.TILE_LAYOUT)
            x_BTC = _all_gather_t(self.ccl_manager, x_BTC, self.parallel_config)
            x_BTC = ttnn.to_layout(x_BTC, ttnn.ROW_MAJOR_LAYOUT)

        # Input C must match the aligned-C the conv weight was allocated for.
        x_BTC = _pad_channels_to_aligned(x_BTC, self.mesh_device)
        x_zs = _zero_stuff_t(x_BTC, stride=self.stride, mesh_device=self.mesh_device)
        x_padded = _zero_pad_t(x_zs, self.external_pad_each, self.external_pad_each, self.mesh_device)
        y = self.conv(x_padded)

        if sharded:
            y = _partition_t(y, self.parallel_config)  # ROW_MAJOR: no tile-aligned offset needed

        if ch_axis is not None:
            # Re-pad C_out to unit so the per-chip C-shard is TILE-legal.
            y = _pad_channels_to_aligned(y, self.mesh_device, channel_align=channel_align_unit(self.parallel_config))
            y = partition_channel(y, self.parallel_config, dim=2)

        return y


class Snake(Module):
    """``y = x + (1 / (α + ε)) · sin(α · x)²``. α has shape ``(1, 1, C)`` per channel."""

    def __init__(
        self,
        channels: int,
        *,
        alpha_logscale: bool = False,
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.float32,
        parallel_config=None,
    ) -> None:
        super().__init__()
        self.channels = channels
        self.alpha_logscale = alpha_logscale  # if True, learned param is log(α)
        self.eps = 1e-9
        self.parallel_config = parallel_config
        # Under channel-TP, pad α to the C-shard unit to match the padded activation.
        unit = channel_align_unit(parallel_config) if channel_axis(parallel_config) is not None else 1
        self._aligned_channels = aligned_channels(channels, unit)
        self.alpha = Parameter(total_shape=[1, 1, self._aligned_channels], device=mesh_device, dtype=dtype)
        self._alpha_shard = None

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        if "alpha" in state:
            a = state["alpha"]
            if self.alpha_logscale:
                a = torch.exp(a)
            a = a.reshape(1, 1, -1)
            if a.shape[-1] < self._aligned_channels:
                a = torch.nn.functional.pad(a, (0, self._aligned_channels - a.shape[-1]))
            state["alpha"] = a.contiguous()

    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
        # α is per-channel; C-shard it under channel-TP.
        if self._alpha_shard is None:
            self._alpha_shard = partition_channel(self.alpha.data, self.parallel_config, dim=2)
        a = self._alpha_shard
        ax = ttnn.multiply(x_BTC, a)
        s = ttnn.sin(ax)
        s2 = ttnn.multiply(s, s)
        inv = ttnn.reciprocal(ttnn.add(a, self.eps))
        return ttnn.add(x_BTC, ttnn.multiply(s2, inv))


class SnakeBeta(Module):
    """``y = x + (1 / (β + ε)) · sin(α · x)²``. α, β both ``(1, 1, C)`` learned."""

    def __init__(
        self,
        channels: int,
        *,
        alpha_logscale: bool = False,
        mesh_device: ttnn.MeshDevice,
        dtype: ttnn.DataType = ttnn.float32,
        parallel_config=None,
    ) -> None:
        super().__init__()
        self.channels = channels
        self.alpha_logscale = alpha_logscale
        self.eps = 1e-9
        self.parallel_config = parallel_config
        # Under channel-TP, pad α/β to the C-shard unit to match the padded activation.
        unit = channel_align_unit(parallel_config) if channel_axis(parallel_config) is not None else 1
        self._aligned_channels = aligned_channels(channels, unit)
        self.alpha = Parameter(total_shape=[1, 1, self._aligned_channels], device=mesh_device, dtype=dtype)
        self.beta = Parameter(total_shape=[1, 1, self._aligned_channels], device=mesh_device, dtype=dtype)
        self._ab_shard = None

    def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None:
        for name in ("alpha", "beta"):
            if name in state:
                t = state[name]
                if self.alpha_logscale:
                    t = torch.exp(t)
                t = t.reshape(1, 1, -1)
                real = t.shape[-1]
                if real < self._aligned_channels:
                    t = torch.nn.functional.pad(t, (0, self._aligned_channels - real))
                if name == "beta":
                    t = t.clone()
                    t[..., :real] = t[..., :real] + self.eps
                    if real < self._aligned_channels:
                        t[..., real:] = 1.0
                state[name] = t.contiguous()

    def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor:
        # α, β per-channel; C-shard under channel-TP.
        if self._ab_shard is None:
            self._ab_shard = (
                partition_channel(self.alpha.data, self.parallel_config, dim=2),
                partition_channel(self.beta.data, self.parallel_config, dim=2),
            )
        a, b = self._ab_shard

        x_tile = ttnn.to_layout(x_BTC, ttnn.TILE_LAYOUT) if x_BTC.layout != ttnn.TILE_LAYOUT else x_BTC
        y = ttnn.snake_beta(x_tile, a, b)
        return y