Functional

Functions for a graph's body, under torch.nn.functional's names.

Each records the operator steps that compute it and returns their handle. Weights are a checkpoint's (out, in), as torch.nn.Linear stores them. A handle's own + - * / are steps too, and @ of one row.

ConvState dataclass

ConvState(read: Handle, write: Handle, window: State, tail: State)

What a short causal conv carries between calls.

read, the K - 1 rows before a call's first, and write, where the K - 1 ending it go; and the scratch states a call stages them in, window (2, chunk, channels) and tail (K - 1, channels).

linear

linear(x: Handle | State, weight, bias=None, *, precision: Precision, **gemv) -> Handle

Return x @ weight.T, plus bias on every row if given.

weight is (out, in), as torch.nn.Linear stores it. One row runs on GEMV, which always rounds once in float32 regardless of precision; several rows run on flm.ProjectionGEMM at precision, or on GEMM at Precision.HIGHEST.

Raises:

Type	Description
ValueError	x of another rank, or GEMV tunables for several rows.

Source code in iron/functional.py

def linear(
    x: Handle | State, weight, bias=None, *, precision: Precision, **gemv
) -> Handle:
    """Return ``x @ weight.T``, plus ``bias`` on every row if given.

    ``weight`` is ``(out, in)``, as ``torch.nn.Linear`` stores it. One row
    runs on GEMV, which always rounds once in float32 regardless of
    ``precision``; several rows run on flm.ProjectionGEMM at ``precision``,
    or on GEMM at ``Precision.HIGHEST``.

    Raises:
        ValueError: ``x`` of another rank, or GEMV tunables for several rows.
    """
    one_row = len(x.shape) == 1 or (len(x.shape) == 2 and x.shape[0] == 1)
    if len(x.shape) not in (1, 2):
        raise ValueError(f"linear(): x {tuple(x.shape)} is not (in,) or (rows, in)")
    if one_row:
        y = GEMV.step(weight, x, **gemv)
    elif gemv:
        raise ValueError(f"linear(): {x.shape[0]} rows take no GEMV {sorted(gemv)}")
    elif precision == Precision.HIGHEST:
        y = GEMM.step(x, weight, b_col_maj=True, precision=Precision.HIGHEST)
    else:
        y = ProjectionGEMM.step(
            x,
            weight,
            b_col_maj=True,
            epilogue_modes=(Epilogue.NONE,),
            precision=precision,
        )
    if bias is None:
        return y
    # A row added to every row, also for GEMV's one-dimensional output.
    rows = math.prod(y.shape) // bias.shape[0]
    return (y.reshape(rows, bias.shape[0]) + bias).reshape(y.shape)

rms_norm

rms_norm(x: Handle, normalized_shape, weight=None, eps: float | None = None) -> Handle

x over the RMS of its last axis, times weight if given.

eps None is x's dtype's machine epsilon, as torch's.

Raises:

Type	Description
ValueError	normalized_shape is not x's last axis.

Source code in iron/functional.py

def rms_norm(
    x: Handle, normalized_shape, weight=None, eps: float | None = None
) -> Handle:
    """``x`` over the RMS of its last axis, times ``weight`` if given.

    ``eps`` None is ``x``'s dtype's machine epsilon, as torch's.

    Raises:
        ValueError: ``normalized_shape`` is not ``x``'s last axis.
    """
    _last_axis("rms_norm", x, normalized_shape)
    if eps is None:
        eps = float(ml_dtypes.finfo(x.dtype).eps)
    return RMSNorm.step(x, weight=weight, epsilon=eps)

layer_norm

layer_norm(x: Handle, normalized_shape, weight=None, bias=None, eps: float = 1e-05) -> Handle

Layer-normalize x over its last axis, as torch's layer_norm.

Subtract the mean, divide by the standard deviation (eps under the root), then multiply by weight and add bias when given.

Raises:

Type	Description
ValueError	normalized_shape is not x's last axis.

Source code in iron/functional.py

def layer_norm(
    x: Handle, normalized_shape, weight=None, bias=None, eps: float = 1e-5
) -> Handle:
    """Layer-normalize ``x`` over its last axis, as torch's ``layer_norm``.

    Subtract the mean, divide by the standard deviation (``eps`` under the
    root), then multiply by ``weight`` and add ``bias`` when given.

    Raises:
        ValueError: ``normalized_shape`` is not ``x``'s last axis.
    """
    _last_axis("layer_norm", x, normalized_shape)
    y = LayerNorm.step(x, weight=weight, epsilon=eps)
    return y if bias is None else y + bias

silu

silu(x: Handle) -> Handle

Apply SiLU, x * sigmoid(x), as torch's.

Source code in iron/functional.py

def silu(x: Handle) -> Handle:
    """Apply SiLU, ``x * sigmoid(x)``, as torch's."""
    return SiLU.step(x)

gelu

gelu(x: Handle, approximate: str = 'none') -> Handle

Apply GELU, exact by default as torch's.

Or by its tanh approximation (approximate="tanh", gelu_pytorch_tanh).

Source code in iron/functional.py

def gelu(x: Handle, approximate: str = "none") -> Handle:
    """Apply GELU, exact by default as torch's.

    Or by its tanh approximation (``approximate="tanh"``,
    ``gelu_pytorch_tanh``).
    """
    return GELU.step(x, approximate=approximate)

rope

rope(x: Handle, angles: Handle) -> Handle

x (rows, head_dim) turned by RoPE, a row of angles per row.

Source code in iron/functional.py

def rope(x: Handle, angles: Handle) -> Handle:
    """``x`` ``(rows, head_dim)`` turned by RoPE, a row of ``angles`` per row."""
    return RoPE.step(x, angles)

causal_conv1d

causal_conv1d(x: Handle, taps, state: ConvState) -> Handle

Return x (rows, channels) through a depth-wise causal filter of K taps.

Each channel its own column of taps (K, channels): output row t is sum_k taps[k] * x[t - K + 1 + k]. The K - 1 rows before x's first row are read from state.read, and the K - 1 rows ending this call (counting those read) are written to state.write for the next call. When x's rows are bounded per call (x[:n]), the call is a prompt chunk of at most chunk rows; otherwise x is one row, a decode step.

Source code in iron/functional.py

def causal_conv1d(x: Handle, taps, state: ConvState) -> Handle:
    """Return ``x`` ``(rows, channels)`` through a depth-wise causal filter of ``K`` taps.

    Each channel its own column of ``taps`` ``(K, channels)``: output row
    ``t`` is ``sum_k taps[k] * x[t - K + 1 + k]``. The ``K - 1``
    rows before ``x``'s first row are read from ``state.read``, and the
    ``K - 1`` rows ending this call (counting those read) are written to
    ``state.write`` for the next call. When ``x``'s rows are bounded per
    call (``x[:n]``), the call is a prompt chunk of at most ``chunk`` rows;
    otherwise ``x`` is one row, a decode step.
    """
    read, write, window, tail = state.read, state.write, state.window, state.tail
    K, C = taps.shape[0], window.shape[1]
    # Block 0 ends with the rows before the call; block 1 is the call's.
    flat = window.reshape(2 * C, x.shape[1])
    first = C - (K - 1)
    flat[first:C] = read
    rows = x.valid.get(0)
    if rows is None:
        flat[C] = x.reshape(x.shape[1])
        y = x * taps[K - 1]
        for k in range(K - 1):
            y = y + flat[first + k : first + k + 1] * taps[k]
        Copy.step(flat[first + 1 : C + 1], write)
        return y
    window[1, :rows] = x
    y = flat[first : first + C] * taps[0]
    for k in range(1, K):
        y = y + flat[first + k : first + k + C] * taps[k]
    y = y[:rows]
    for j in range(K - 1):
        tail[j] = flat[first + rows + j]
    Copy.step(tail, write)
    return y