Temporal Convolutional Network

"""A dilated causal convolution stack - transcribed.

Two modifications turn an ordinary convolution into a sequence model. Causal:
only look left. Dilated: skip a growing gap between taps, so the receptive
field doubles per layer while the parameter count grows linearly.
"""


def receptive_field(layers, kernel):
    """R = 1 + sum over layers of (k - 1) * dilation, dilations doubling.

    This is arithmetic, not a hyperparameter. Compute it and check it covers
    the longest dependency in the problem BEFORE training anything - a
    dependency one step beyond it is structurally invisible, not merely hard.
    """
    total = 1
    for layer in range(layers):
        total += (kernel - 1) * (2 ** layer)
    return total


def causal_conv1d(sequence, weights, bias, dilation):
    """One dilated causal convolution over a single-channel sequence.

    sequence[t] contributes only to outputs at positions >= t. The loop below
    reads BACKWARDS from t, never forwards, which is what makes it causal -
    and the reason a symmetric padding scheme silently breaks the model.
    """
    kernel = len(weights)
    out = [0.0] * len(sequence)

    for t in range(len(sequence)):
        total = bias
        for i in range(kernel):
            source = t - i * dilation
            # Positions before the start are treated as zero: this is the
            # LEFT padding, and it must never be applied on the right.
            if source >= 0:
                total += weights[i] * sequence[source]
        out[t] = total

    return out


def relu(values):
    return [value if value > 0.0 else 0.0 for value in values]


def residual_block(sequence, weights_a, bias_a, weights_b, bias_b, dilation):
    """Two dilated convolutions plus a skip connection.

    The residual is the same idea gating uses in a recurrent model: an
    ADDITIVE path the gradient travels without being multiplied. Here it is
    applied to depth rather than to time, and it is what makes a deep stack
    trainable.
    """
    hidden = relu(causal_conv1d(sequence, weights_a, bias_a, dilation))
    output = causal_conv1d(hidden, weights_b, bias_b, dilation)
    return [output[t] + sequence[t] for t in range(len(sequence))]


def forward(sequence, blocks):
    """blocks is a list of (weights_a, bias_a, weights_b, bias_b) tuples.

    Dilation doubles per block, so a stack of L blocks reaches back
    exponentially far using only L * 2 * k weights.
    """
    current = list(sequence)
    for layer, (w_a, b_a, w_b, b_b) in enumerate(blocks):
        current = residual_block(current, w_a, b_a, w_b, b_b, 2 ** layer)
    return current

"""A TCN - typed, multi-channel, causal padding verified, as one matmul."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


@dataclass(frozen=True)
class DilatedBlock:
    """One residual block: two dilated causal convolutions and a skip."""

    weights_a: Tensor    # (kernel, in_channels, out_channels)
    bias_a: Vector
    weights_b: Tensor
    bias_b: Vector
    dilation: int

    @property
    def kernel(self) -> int:
        return self.weights_a.shape[0]


def receptive_field(blocks: list[DilatedBlock]) -> int:
    """R = 1 + sum over layers of (k - 1) * dilation, counting BOTH
    convolutions in each residual block.

    Arithmetic, not a hyperparameter. A dependency one step beyond R is
    structurally invisible rather than merely hard to learn, so this should be
    computed and checked against the problem before training anything.
    """
    return 1 + sum(2 * (block.kernel - 1) * block.dilation for block in blocks)


def causal_conv1d(sequence: Tensor, weights: Tensor, bias: Vector, dilation: int) -> Tensor:
    """Dilated causal convolution over (batch, time, channels).

    Implemented as an explicit gather plus one matrix product rather than a
    loop over positions: the k taps are gathered into a (batch, time, k *
    in_channels) block, which is then one GEMM against the flattened kernel.
    That is how every framework implements a convolution, and it is why a
    convolution is fast at all.
    """
    if sequence.ndim != 3:
        raise ValueError(f"expected (batch, time, channels), got {sequence.shape}")

    batch, steps, in_channels = sequence.shape
    kernel = weights.shape[0]

    # LEFT padding only. Padding both sides leaks the future and produces a
    # model that validates beautifully and fails in production - the most
    # damaging and most common error with this architecture.
    padding = (kernel - 1) * dilation
    padded = np.pad(sequence, ((0, 0), (padding, 0), (0, 0)))

    # Gather the k dilated taps for every output position.
    gathered = np.empty((batch, steps, kernel * in_channels), dtype=np.float64)
    for i in range(kernel):
        start = padding - i * dilation
        gathered[:, :, i * in_channels : (i + 1) * in_channels] = padded[
            :, start : start + steps
        ]

    flat_weights = weights.reshape(kernel * in_channels, -1)
    return gathered @ flat_weights + bias


def forward(sequence: Tensor, blocks: list[DilatedBlock]) -> Tensor:
    """Run the stack. Every position is computed independently, so the whole
    sequence is processed in one pass - the property that distinguishes this
    from a recurrent model."""
    current = sequence

    for block in blocks:
        hidden = np.maximum(
            causal_conv1d(current, block.weights_a, block.bias_a, block.dilation), 0.0
        )
        output = causal_conv1d(hidden, block.weights_b, block.bias_b, block.dilation)

        # The residual: an ADDITIVE path the gradient travels without being
        # multiplied - the same idea gating uses, applied to depth rather
        # than time, and what makes a deep stack trainable.
        if output.shape[2] != current.shape[2]:
            raise ValueError(
                f"residual shape mismatch: block outputs {output.shape[2]} channels "
                f"but input has {current.shape[2]}; a 1x1 projection is needed"
            )
        current = output + current

    return current


def assert_causal(blocks: list[DilatedBlock], steps: int, channels: int) -> None:
    """Verify causality directly rather than assuming it.

    Perturb the LAST input position and confirm no earlier output changed. A
    padding or shift mistake is silent - it produces excellent validation
    numbers - so this check belongs in the test suite rather than in a
    reviewer's head.
    """
    baseline = np.zeros((1, steps, channels), dtype=np.float64)
    perturbed = baseline.copy()
    perturbed[0, -1] = 1.0

    before = forward(baseline, blocks)
    after = forward(perturbed, blocks)

    if not np.allclose(before[0, :-1], after[0, :-1]):
        raise AssertionError("causality violated: a past output changed with a future input")

"""A TCN - preallocated gather buffers, cached activations for streaming."""

import numpy as np
from numpy.typing import NDArray

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Tensor = NDArray[np.float64]


class StreamingTcn:
    """Two changes, for the two regimes this architecture is used in.

    Batch: the gather buffer for each layer is allocated once and rewritten,
    rather than allocating a (batch, time, k * channels) block per convolution
    per forward pass. On a deep stack over a long sequence that is the
    dominant allocation, and activation memory is what a TCN exhausts rather
    than time.

    Streaming: naive autoregressive generation recomputes the entire receptive
    field for every output, which is the single largest performance mistake
    with this architecture. Keeping a small ring buffer per layer - just the
    k dilated taps that layer needs - makes each new step O(L * k * c^2)
    instead of O(R * L * k * c^2).
    """

    def __init__(self, blocks: list[tuple[Tensor, Vector, Tensor, Vector, int]],
                 batch: int, max_steps: int) -> None:
        self._blocks = blocks
        # One contiguous float32 buffer per layer, sized once. float32 halves
        # the traffic on the activation stack, which is the binding resource.
        self._gathered: list[Matrix] = []
        self._rings: list[Matrix] = []

        for weights_a, _, _, _, dilation in blocks:
            kernel, in_channels, _ = weights_a.shape
            self._gathered.append(
                np.zeros((batch, max_steps, kernel * in_channels), dtype=np.float32)
            )
            # Ring buffer holding only the taps this layer reads: k positions
            # spaced by the dilation, not the whole receptive field.
            self._rings.append(
                np.zeros((batch, (kernel - 1) * dilation + 1, in_channels), dtype=np.float32)
            )

    def _conv(self, sequence: Tensor, weights: Tensor, bias: Vector, dilation: int,
              buffer: Matrix) -> Tensor:
        batch, steps, in_channels = sequence.shape
        kernel = weights.shape[0]
        padding = (kernel - 1) * dilation

        padded = np.zeros((batch, steps + padding, in_channels), dtype=np.float32)
        padded[:, padding:] = sequence      # LEFT padding only

        gathered = buffer[:batch, :steps]
        for i in range(kernel):
            start = padding - i * dilation
            # Written into the preallocated buffer rather than a new array.
            np.copyto(
                gathered[:, :, i * in_channels : (i + 1) * in_channels],
                padded[:, start : start + steps],
            )

        return gathered @ weights.reshape(kernel * in_channels, -1) + bias

    def forward(self, sequence: Tensor) -> Tensor:
        """Whole-sequence pass. Every position computes independently, which
        is what makes this fully parallel and a recurrent model not."""
        current = np.ascontiguousarray(sequence, dtype=np.float32)

        for index, (weights_a, bias_a, weights_b, bias_b, dilation) in enumerate(self._blocks):
            hidden = self._conv(current, weights_a, bias_a, dilation, self._gathered[index])
            np.maximum(hidden, 0.0, out=hidden)
            output = self._conv(hidden, weights_b, bias_b, dilation, self._gathered[index])
            current = output + current      # residual

        return current

    def step(self, x: Matrix) -> Matrix:
        """One new timestep, using cached taps.

        This is the difference between viable and unusable autoregressive
        generation: each layer keeps only the k positions it actually reads,
        so a new output costs one pass through the stack rather than a
        recomputation of the whole receptive field.
        """
        current = x

        for index, (weights_a, bias_a, weights_b, bias_b, dilation) in enumerate(self._blocks):
            ring = self._rings[index]
            # Shift in the new value; the ring holds exactly the span this
            # layer's dilated taps reach across, and nothing more.
            ring[:, :-1] = ring[:, 1:]
            ring[:, -1] = current

            kernel = weights_a.shape[0]
            taps = np.concatenate(
                [ring[:, -1 - i * dilation] for i in range(kernel)], axis=1
            )

            hidden = np.maximum(taps @ weights_a.reshape(-1, weights_a.shape[2]) + bias_a, 0.0)
            output = hidden @ weights_b.reshape(-1, weights_b.shape[2]) + bias_b
            current = output + current

        return current