Language models

iron.lm is what every language model shares; each model is a package under iron/applications/ over it. The agent instructions describe how the modules fit together.

config

A model's shape.

Config, read from a checkpoint's own config.json, and a decoder's, DecoderConfig, with its RoPE table (rope_angles).

Config dataclass

Config(*, eps: float = 1e-06, precision: Precision)

A model's shape, built from a checkpoint's config.json (from_checkpoint).

precision is every projection's, chosen per application on its measured accuracy, with no default here.

from_checkpoint classmethod

from_checkpoint(path: str | Path) -> Self

cls's fields read from path's own config.json.

Each at its metadata["json"] path (under json_prefix), converted by metadata["convert"] if it names one.

Parameters:

Name	Type	Description	Default
path	str | Path	A config.json, or a checkpoint's directory or file (its directory's config.json is read).	required

Raises:

Type	Description
FileNotFoundError	No config.json is there.
ValueError	The file lacks a field that has no default.

Source code in iron/lm/config.py

@classmethod
def from_checkpoint(cls, path: str | Path) -> Self:
    """``cls``'s fields read from ``path``'s own ``config.json``.

    Each at its ``metadata["json"]`` path (under ``json_prefix``),
    converted by ``metadata["convert"]`` if it names one.

    Args:
        path: A ``config.json``, or a checkpoint's directory or file
            (its directory's ``config.json`` is read).

    Raises:
        FileNotFoundError: No ``config.json`` is there.
        ValueError: The file lacks a field that has no default.
    """
    path = Path(path)
    if path.name != "config.json":
        path = (path if path.is_dir() else path.parent) / "config.json"
    return cls.from_json(json.loads(path.read_text()), path)

from_json classmethod

from_json(data: dict, source: str | Path = 'config.json') -> Self

from_checkpoint on data, a config.json already parsed.

source names it in an error.

Raises:

Type	Description
ValueError	data lacks a field that has no default, or cls states no default for a field config.json does not set (an application's choice, as precision).

Source code in iron/lm/config.py

@classmethod
def from_json(cls, data: dict, source: str | Path = "config.json") -> Self:
    """``from_checkpoint`` on ``data``, a ``config.json`` already parsed.

    ``source`` names it in an error.

    Raises:
        ValueError: ``data`` lacks a field that has no default, or
            ``cls`` states no default for a field ``config.json`` does
            not set (an application's choice, as ``precision``).
    """
    kwargs, missing = {}, []
    for f in dataclasses.fields(cls):
        key = cls.json_key(f)
        if key is None:
            required = dataclasses.MISSING
            if f.default is required and f.default_factory is required:
                raise ValueError(f"{cls.__name__} states no {f.name}")
            continue
        value = _dotted(data, key)
        # Hugging Face writes an unset option as null.
        if value is _MISSING or value is None:
            required = dataclasses.MISSING
            if f.default is required and f.default_factory is required:
                missing.append(key)
            continue
        convert = f.metadata.get("convert")
        kwargs[f.name] = convert(value) if convert else value
    if missing:
        raise ValueError(f"{source} lacks {missing}")
    return cls(**kwargs)

json_key classmethod

json_key(f: Field) -> str | None

Where config.json holds field f; None: nowhere.

Source code in iron/lm/config.py

@classmethod
def json_key(cls, f: dataclasses.Field) -> str | None:
    """Where ``config.json`` holds field ``f``; None: nowhere."""
    key = f.metadata.get("json", f.name)
    if key is None or key.startswith("/"):
        return None if key is None else key[1:]
    return cls.json_prefix + key

to_json

to_json() -> dict

Return the fields as config.json names them, from_checkpoint's inverse.

A nested Config gives its own to_json(), another dataclass its fields. A not_json field, or one that is None or not JSON, is left out.

Source code in iron/lm/config.py

def to_json(self) -> dict:
    """Return the fields as ``config.json`` names them, ``from_checkpoint``'s inverse.

    A nested ``Config`` gives its own ``to_json()``, another dataclass
    its fields. A ``not_json`` field, or one that is None or not JSON,
    is left out.
    """
    out = {}
    for f in dataclasses.fields(self):
        key = self.json_key(f)
        value = getattr(self, f.name)
        if isinstance(value, Config):
            value = value.to_json()
        elif dataclasses.is_dataclass(value) and not isinstance(value, type):
            value = dataclasses.asdict(value)
        if key is None or not isinstance(
            value, (bool, int, float, str, tuple, dict)
        ):
            continue
        *path, leaf = key.split(".")
        node = out
        for part in path:
            node = node.setdefault(part, {})
        node[leaf] = list(value) if isinstance(value, tuple) else value
    return out

DecoderConfig dataclass

DecoderConfig(vocab_size: int, hidden_size: int, num_hidden_layers: int, num_attention_heads: int, num_key_value_heads: int, intermediate_size: int, rope_theta: float | None, max_seq_len: int, rope_scaling: RopeScaling | None = None, rotary_dim: int | None = not_json(None), rope_attention_scaling: float = not_json(1.0), prefill_chunk: int = not_json(2048), hidden_act: str = 'silu', tie_word_embeddings: bool = True, *, eps: float = 1e-06, precision: Precision, head_dim: int | None = None, attention_bias: bool = False, mlp_bias: bool = False, sliding_window: int | None = None, use_sliding_window: bool | None = None)

Bases: Config

A decoder's shape, named as a Hugging Face checkpoint's own config.json does where it has one.

max_seq_len is the checkpoint's native context unless a caller lowers it.

window property

window: int | None

The positions before a query each layer attends to (KVCache's window), sliding_window less the query's own.

None is all.

rope_width property

rope_width: int

The angle table's own width: head_dim, or rotary_dim under partial rotary.

angles

angles() -> ndarray | None

Return the RoPE table, a row per position; None with no rope_theta.

Source code in iron/lm/config.py

def angles(self) -> np.ndarray | None:
    """Return the RoPE table, a row per position; None with no ``rope_theta``."""
    if self.rope_theta is None:
        return None
    angles = rope_angles(
        self.rope_width, self.max_seq_len, self.rope_theta, self.rope_scaling
    )
    if self.rope_attention_scaling != 1.0:
        angles = angles * np.float32(self.rope_attention_scaling)
    return angles

rope_angles

rope_angles(head_dim: int, context_length: int, rope_theta: float = 500000.0, scaling: RopeScaling | None = None) -> ndarray

Return the RoPE table, (context_length, head_dim) float32, cos and sin interleaved.

Each transcendental is evaluated in float64 and rounded once.

Source code in iron/lm/config.py

def rope_angles(
    head_dim: int,
    context_length: int,
    rope_theta: float = 500000.0,
    scaling: RopeScaling | None = None,
) -> np.ndarray:
    """Return the RoPE table, ``(context_length, head_dim)`` float32, cos and sin interleaved.

    Each transcendental is evaluated in float64 and rounded once.
    """
    exponents = np.arange(0, head_dim, 2, dtype=np.float32) / np.float32(head_dim)
    inv_freq = 1.0 / np.power(rope_theta, exponents.astype(np.float64))
    if scaling is not None:
        inv_freq = scaling(inv_freq)
    inv_freq = inv_freq.astype(np.float32)
    freqs = np.outer(np.arange(context_length, dtype=np.float32), inv_freq)
    angles = np.empty((context_length, head_dim), dtype=np.float32)
    angles[:, ::2] = np.cos(freqs.astype(np.float64))
    angles[:, 1::2] = np.sin(freqs.astype(np.float64))
    return angles

unscaled

unscaled(rope_scaling: dict) -> None

Return None for a checkpoint's rope_scaling that leaves RoPE as it is.

Raises:

Type	Description
ValueError	It scales RoPE, which this config builds no table for.

Source code in iron/lm/config.py

def unscaled(rope_scaling: dict) -> None:
    """Return None for a checkpoint's ``rope_scaling`` that leaves RoPE as it is.

    Raises:
        ValueError: It scales RoPE, which this config builds no table for.
    """
    kind = rope_scaling.get("rope_type", rope_scaling.get("type"))
    if kind != "default":
        raise ValueError(f"rope_scaling {rope_scaling}: no RoPE table here builds it")

not_json

not_json(default, **kwargs) -> Any

Declare a field config.json does not set, an IRON choice.

Kept at default by from_checkpoint and left out by to_json.

Source code in iron/lm/config.py

def not_json(default, **kwargs) -> Any:
    """Declare a field ``config.json`` does not set, an IRON choice.

    Kept at ``default`` by ``from_checkpoint`` and left out by ``to_json``.
    """
    return dataclasses.field(default=default, metadata={"json": None}, **kwargs)

checkpoint

A model's weights: Checkpoint.tree(), the tensors by their own names.

model.layers.3.self_attn.q_proj.weight is weights.model.layers[3].self_attn.q_proj.weight.

Checkpoint

Checkpoint(path)

A .safetensors file or a PyTorch checkpoint.

Or, given its .safetensors.index.json, every shard a sharded checkpoint splits across, its tensors merged into one mapping, tensors.

Source code in iron/lm/checkpoint.py

def __init__(self, path):
    self.path = Path(path)
    weight_map = None
    paths = [self.path]
    if self.path.suffix == ".json":
        weight_map = json.loads(self.path.read_text())["weight_map"]
        names = sorted(set(weight_map.values()))
        for n in names:
            if Path(n).name != n:
                raise ValueError(f"{self.path}: {n} is not a file beside the index")
        paths = [self.path.parent / n for n in names]
    self._shards = [_Shard(p) for p in paths]
    self.tensors: dict[str, np.ndarray] = {}
    for shard in self._shards:
        if weight_map is not None:
            named = {n for n, s in weight_map.items() if s == shard.path.name}
            if named != shard.tensors.keys():
                raise ValueError(
                    f"{shard.path}: holds "
                    f"{sorted(shard.tensors.keys() - named)} the index does "
                    f"not name and lacks {sorted(named - shard.tensors.keys())}"
                )
        self.tensors.update(shard.tensors)

release

release(array: ndarray) -> None

Drop the process's pages of array if it is a view of one of the files.

Anything else is left alone.

Source code in iron/lm/checkpoint.py

def release(self, array: np.ndarray) -> None:
    """Drop the process's pages of ``array`` if it is a view of one of the files.

    Anything else is left alone.
    """
    for shard in self._shards:
        if shard.release(array):
            return

fingerprint

fingerprint() -> str

Return a digest of each tensor's name, type, shape and first/last 4 KiB.

What tells a fine-tune from its base without reading either whole.

Source code in iron/lm/checkpoint.py

def fingerprint(self) -> str:
    """Return a digest of each tensor's name, type, shape and first/last 4 KiB.

    What tells a fine-tune from its base without reading either whole.
    """
    ends = []
    for name, t in self.tensors.items():
        n = 4096 // t.itemsize
        ends.append([name, f"{t.dtype}{t.shape}", t.flat[:n], t.flat[-n:]])
    return fingerprint(ends)

tree

tree(fill: Literal['read', 'random', 'unread'] = 'read', seed: int = 0) -> SimpleNamespace

Return the checkpoint's tensors as a tree of param, split at the dots.

model.layers.3.weight becomes tree.model.layers[3].weight.

Parameters:

Name	Type	Description	Default
fill	Literal['read', 'random', 'unread']	What the arrays hold: "read" the tensors (lazy views); "random" or "unread" filled's, which reads nothing past the header.	'read'
seed	int	The seed "random" draws at.	0

Raises:

Type	Description
ValueError	fill is none of these.

Source code in iron/lm/checkpoint.py

def tree(
    self, fill: Literal["read", "random", "unread"] = "read", seed: int = 0
) -> SimpleNamespace:
    """Return the checkpoint's tensors as a tree of ``param``, split at the dots.

    ``model.layers.3.weight`` becomes ``tree.model.layers[3].weight``.

    Args:
        fill: What the arrays hold: ``"read"`` the tensors (lazy views);
            ``"random"`` or ``"unread"`` ``filled``'s, which reads
            nothing past the header.
        seed: The seed ``"random"`` draws at.

    Raises:
        ValueError: ``fill`` is none of these.
    """
    if fill not in ("read", "random", "unread"):
        raise ValueError(f"fill is read, random or unread, not {fill!r}")
    tensors = self.tensors
    if fill != "read":
        named = {n: (t.shape, t.dtype) for n, t in tensors.items()}
        tensors = filled(fill, named, seed)
    return nest({n: param(t) for n, t in tensors.items()}, file=self.path)

weights_file

weights_file(directory) -> Path

Return a checkpoint directory's weights file.

A sharded checkpoint's model.safetensors.index.json, else model.safetensors.

Source code in iron/lm/checkpoint.py

def weights_file(directory) -> Path:
    """Return a checkpoint directory's weights file.

    A sharded checkpoint's ``model.safetensors.index.json``, else
    ``model.safetensors``.
    """
    index = Path(directory) / "model.safetensors.index.json"
    return index if index.exists() else index.with_name("model.safetensors")

leaves

leaves(tree, prefix: str = '') -> Iterator[tuple[str, Any]]

Return each leaf of tree, a Checkpoint.tree()'s namespaces and lists.

By its dotted name: nest's inverse.

Source code in iron/lm/checkpoint.py

def leaves(tree, prefix: str = "") -> Iterator[tuple[str, Any]]:
    """Return each leaf of ``tree``, a ``Checkpoint.tree()``'s namespaces and lists.

    By its dotted name: ``nest``'s inverse.
    """
    if isinstance(tree, (list, SimpleNamespace)):
        items = enumerate(tree) if isinstance(tree, list) else vars(tree).items()
        for name, value in items:
            yield from leaves(value, f"{prefix}{name}.")
    else:
        yield prefix[:-1], tree

nest

nest(named: Mapping[str, Any], path: str = '', file: str | Path = _FILE) -> SimpleNamespace

Nest values by dotted name, as Checkpoint.tree() nests them.

Under path (a dotted prefix) of file, which a missing name names.

Source code in iron/lm/checkpoint.py

def nest(
    named: Mapping[str, Any], path: str = "", file: str | Path = _FILE
) -> SimpleNamespace:
    """Nest values by dotted name, as ``Checkpoint.tree()`` nests them.

    Under ``path`` (a dotted prefix) of ``file``, which a missing name
    names.
    """
    nested: dict = {}
    for name, value in named.items():
        *parents, leaf = name.split(".")
        node = nested
        for key in parents:
            node = node.setdefault(key, {})
        node[leaf] = value
    return _tree(nested, path, file)

subtree

subtree(tree, **modules)

Build a tree of modules, each a module of tree, at their names.

A checkpoint's part as a model loads it (a text decoder of a vision-language checkpoint), a missing name naming tree's file.

Source code in iron/lm/checkpoint.py

def subtree(tree, **modules):
    """Build a tree of ``modules``, each a module of ``tree``, at their names.

    A checkpoint's part as a model loads it (a text decoder of a
    vision-language checkpoint), a missing name naming ``tree``'s file.
    """
    file = tree._file if isinstance(tree, (_Module, _Layers)) else _FILE
    return _tree(modules, "", file)

renamed

renamed(tree, renames: Mapping[str, str])

Copy tree with each leaf under a run of whole dotted names moved.

A leaf among renames' keys moves to the run its value names, the rest stay where they are; its namespaces are new, so a caller may change them and leave tree as it was.

Source code in iron/lm/checkpoint.py

def renamed(tree, renames: Mapping[str, str]):
    """Copy ``tree`` with each leaf under a run of whole dotted names moved.

    A leaf among ``renames``' keys moves to the run its value names, the
    rest stay where they are; its namespaces are new, so a caller may
    change them and leave ``tree`` as it was.
    """

    def moved(path):
        for old, new in renames.items():
            at = f".{path}.".find(f".{old}.")
            if at >= 0:
                return path[:at] + new + path[at + len(old) :]
        return path

    held = isinstance(tree, (_Module, _Layers))
    return nest(
        {moved(path): value for path, value in leaves(tree)},
        tree._path if held else "",
        tree._file if held else _FILE,
    )

blocks

The blocks a model builds from operators, for a graph's body.

A graph here is called in a model's body as it is on its own (Graph): its steps record into the model's. Weights are a checkpoint's (out, in), projected by F.linear.

LogMel

LogMel(frame: int = 320, hop: int = 160, fft: int = 512, mels: int = 128, sample_rate: int = 16000, max_hz: float = 8000.0, mel_floor: float = 0.001)

Bases: Graph

The log-mel spectrogram of float32 samples, as Hugging Face computes it.

Gemma 4's feature extractor, via bf16-limb GEMMs (Limbs) for exact float32 rfft and filterbank products.

Source code in iron/lm/blocks.py

def __init__(
    self,
    frame: int = 320,
    hop: int = 160,
    fft: int = 512,
    mels: int = 128,
    sample_rate: int = 16000,
    max_hz: float = 8000.0,
    mel_floor: float = 1e-3,
):
    if fft < frame:
        raise ValueError(
            f"LogMel: an rfft of {fft} points is shorter than a frame of {frame}"
        )
    self.frame, self.hop, self.fft, self.mels = frame, hop, fft, mels
    self.mel_floor = mel_floor
    # The samples a frame is read as: six limb planes of them are the
    # rfft GEMM's K, whole 64s; the rfft's rows past the frame are zero.
    self.span = -(-frame // 32) * 32
    filters = filterbank(fft, mels, sample_rate, max_hz)
    bins = fft // 2 + 1
    width = -(-bins // 64) * 64
    angle = 2 * np.pi * np.arange(frame)[:, None] * np.arange(bins) / fft
    window = np.hanning(frame + 1)[:-1, None].astype(np.float32)
    # The rfft as a matrix, the window folded in: the real parts, then
    # the imaginary, each padded to the width.
    dft = np.zeros((self.span, 2 * width))
    dft[:frame, :bins] = window * np.cos(angle)
    dft[:frame, width : width + bins] = -window * np.sin(angle)
    self.dft = iron.param(stacked(dft))
    self.filterbank = iron.param(
        stacked(np.pad(filters, ((0, width - bins), (0, 0))))
    )

ImagePatches

ImagePatches(patches: int, image: tuple[int, int], pool: int = 3, positions: int = 10240, patch_columns: int = 8, chunk: int = 4096)

Bases: Graph

A decoded image resized as torchvision's antialiased bicubic resize does it.

Rescaled to [0, 1] and cut row-major into a vision tower's patches rows of 16x16-pixel patches.

Source code in iron/lm/blocks.py

def __init__(
    self,
    patches: int,
    image: tuple[int, int],
    pool: int = 3,
    positions: int = 10240,
    patch_columns: int = 8,
    chunk: int = 4096,
):
    self.patches, self.image, self.pool, self.chunk = patches, image, pool, chunk
    k, side = pool, pool * SIDE
    cores = 2 * require_device().cols
    # The most pooling windows across: Resample's cores hold
    # `patch_columns` patch columns each, one a zero column past the image.
    self.across = (patch_columns * cores - 1) // k
    # Every grid of whole pooling windows the rows hold; the fewest patch
    # columns give the most bands and padded raster rows.
    grids = [(patches // s // k * k, s) for s in range(k, patches // k + 1, k)]
    shared: dict[str, Any] = dict(
        height=side,
        width=side,
        out_height=side,
        out_width=side,
        chunk=chunk,
        patch_columns=patch_columns,
        width_chunks=patch_columns * cores,
        height_chunks=16 * -(-grids[0][0] // 16),
        patch_rows=max(b * (s // cores + 1) * cores for b, s in grids),
        num_aie_columns=cores // 2,
        num_channels=2,
    )
    a, b = image
    self.resample = Resample(
        image_chunks=max(a * -(-3 * b // chunk), b * -(-3 * a // chunk)),
        **shared,
    )
    self.raster = iron.state((shared["patch_rows"], 3 * SIDE**2))
    taps: dict[str, Any] = dict(
        in_size=side,
        out_size=side,
        words=self.resample.words,
        slots=SIDE,
        num_aie_columns=cores // 2,
        num_channels=2,
    )
    self.taps_w = ResampleTaps(chunks=shared["width_chunks"], **taps)
    self.taps_h = ResampleTaps(chunks=shared["height_chunks"], **taps)
    self.patch_positions = PatchPositions(
        rows=patches,
        out_height=side,
        out_width=side,
        pool=k,
        positions=positions,
        cores=cores,
    )

pack

pack(image: ndarray, out_height: int, out_width: int) -> tuple[ndarray, dict]

Build a call's image buffer and per-call sizes for a decoded image.

Resized to out_height by out_width.

Parameters:

Name	Type	Description	Default
image	ndarray	(height, width, 3) uint8.	required
out_height	int	The height to resize to.	required
out_width	int	The width to resize to.	required

Raises:

Type	Description
ValueError	The image or its patches do not fit a call, or Resample cannot resize it.

Source code in iron/lm/blocks.py

def pack(
    self, image: np.ndarray, out_height: int, out_width: int
) -> tuple[np.ndarray, dict]:
    """Build a call's image buffer and per-call sizes for a decoded ``image``.

    Resized to ``out_height`` by ``out_width``.

    Args:
        image: ``(height, width, 3)`` uint8.
        out_height: The height to resize to.
        out_width: The width to resize to.

    Raises:
        ValueError: The image or its patches do not fit a call, or
            ``Resample`` cannot resize it.
    """
    image = np.asarray(image)
    if image.dtype != np.uint8 or image.ndim != 3 or image.shape[2] != 3:
        raise ValueError(
            f"an image is (height, width, 3) uint8, not {image.shape} {image.dtype}"
        )
    H, W, _ = image.shape
    n = out_height // SIDE * (out_width // SIDE)
    line = -(-3 * W // self.chunk) * self.chunk
    op = self.resample
    if n > self.patches or op.image_chunks * self.chunk < H * line:
        a, b = self.image
        raise ValueError(
            f"a {H}x{W} image of {n} patches does not fit a call's "
            f"{self.patches} patches or {a}x{b} image"
        )
    # Resample's own checks refuse a size it cannot resize.
    op.replace(height=H, width=W, out_height=out_height, out_width=out_width)
    buffer = np.zeros((op.image_chunks, self.chunk), np.uint8)
    rows = buffer.reshape(-1)[: H * line].reshape(H, line)
    rows[:, : 3 * W] = image.reshape(H, 3 * W)
    return buffer, dict(
        height=H, width=W, out_height=out_height, out_width=out_width
    )

Weights

Weights(**shapes: tuple[int, ...])

Bases: Graph

Weights of these shapes, by name.

Read where the checkpoint holds them and computed with elsewhere.

Source code in iron/lm/blocks.py

def __init__(self, **shapes: tuple[int, ...]):
    for name, shape in shapes.items():
        setattr(self, name, iron.param(*shape))

Linear

Linear(in_features: int, out_features: int, bias: bool = True, *, precision: Precision)

Bases: Graph

torch.nn.Linear: x @ weight.T, plus bias where it holds one.

At precision (F.linear's, the caller's choice).

Source code in iron/lm/blocks.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    bias: bool = True,
    *,
    precision: Precision,
):
    self.in_features, self.out_features = in_features, out_features
    self.weight = iron.param(out_features, in_features)
    self.bias = iron.param(out_features) if bias else None
    self.precision = precision

split

split(x, *rows: int, **gemv) -> list

Project x through each run of rows of the weight in turn.

Each a fused projection's output; splitting the output columns instead would need a strided view no bounded-row descriptor fits.

Source code in iron/lm/blocks.py

def split(self, x, *rows: int, **gemv) -> list:
    """Project ``x`` through each run of ``rows`` of the weight in turn.

    Each a fused projection's output; splitting the output columns
    instead would need a strided view no bounded-row descriptor fits.
    """
    ends = np.cumsum(rows)
    return [
        F.linear(
            x,
            self.weight[end - n : end],
            None if self.bias is None else self.bias[end - n : end],
            precision=self.precision,
            **gemv,
        )
        for n, end in zip(rows, ends)
    ]

Embedding

Embedding(num_embeddings: int, embedding_dim: int)

Bases: Graph

torch.nn.Embedding: the rows of weight that ids name.

Source code in iron/lm/blocks.py

def __init__(self, num_embeddings: int, embedding_dim: int):
    self.num_embeddings, self.embedding_dim = num_embeddings, embedding_dim
    self.weight = iron.param(num_embeddings, embedding_dim)

RMSNorm

RMSNorm(normalized_shape: int | tuple[int], eps: float | None = None)

Bases: Graph

torch.nn.RMSNorm over the last axis, times weight.

eps None is the input dtype's machine epsilon.

Source code in iron/lm/blocks.py

def __init__(self, normalized_shape: int | tuple[int], eps: float | None = None):
    n = normalized_shape
    self.normalized_shape = (n,) if isinstance(n, int) else tuple(n)
    self.eps = eps
    self.weight = iron.param(*self.normalized_shape)

plus_one staticmethod

plus_one(module) -> None

Make every RMSNorm of module, loaded, scale by 1 + weight.

As Gemma's and Qwen3.5's do, the sum in float32 rounded to bf16 (Hugging Face keeps it float32).

Source code in iron/lm/blocks.py

@staticmethod
def plus_one(module) -> None:
    """Make every ``RMSNorm`` of ``module``, loaded, scale by ``1 + weight``.

    As Gemma's and Qwen3.5's do, the sum in float32 rounded to bf16
    (Hugging Face keeps it float32).
    """
    for _, weight, norm, _ in list(walk(module)):
        if isinstance(norm, RMSNorm):
            norm.weight = iron.param(
                (weight.astype(np.float32) + 1).astype(bfloat16)
            )

LayerNorm

LayerNorm(normalized_shape: int | tuple[int], eps: float = 1e-05)

Bases: Graph

torch.nn.LayerNorm over the last axis, times weight plus bias.

Computed by F.layer_norm.

Source code in iron/lm/blocks.py

def __init__(self, normalized_shape: int | tuple[int], eps: float = 1e-5):
    n = normalized_shape
    self.normalized_shape = (n,) if isinstance(n, int) else tuple(n)
    self.eps = eps
    self.weight = iron.param(*self.normalized_shape)
    self.bias = iron.param(*self.normalized_shape)

Rope

Rope(c, q_scale: ndarray | None = None)

Bases: Graph

A decoder's RoPE table, one row per position (DecoderConfig.angles).

q and k share it unless q_scale gives q a scaled table of its own. A model copies each call's rows once (rows) for every layer to read.

Source code in iron/lm/blocks.py

def __init__(self, c, q_scale: np.ndarray | None = None):
    angles = c.angles()
    self.table = iron.weight(angles.astype(bfloat16))
    self.q_table = None
    if q_scale is not None:
        q = angles * np.asarray(q_scale, np.float32)[:, None]
        self.q_table = iron.weight(q.astype(bfloat16))
    self.chunk, self.width = c.prefill_chunk, c.rope_width

rows

rows(call) -> tuple

Return call's rows of q's table and of k's.

The same rows where they share one: a prompt chunk's from its start, else the decode step's at its position.

Source code in iron/lm/blocks.py

def rows(self, call) -> tuple:
    """Return ``call``'s rows of q's table and of k's.

    The same rows where they share one: a prompt chunk's from its
    ``start``, else the decode step's at its ``position``.
    """

    def read(table):
        if call.prompt:
            start = call.start
            return Copy.step(table[start : start + self.chunk][: call.rows])
        return Copy.step(table[call.position]).reshape(1, self.width)

    k = read(self.table)
    return (k if self.q_table is None else read(self.q_table)), k

table_start

table_start(n: int) -> int

Return the row of the table a prompt of n tokens starts reading from.

0, the table's first.

Source code in iron/lm/blocks.py

def table_start(self, n: int) -> int:
    """Return the row of the table a prompt of ``n`` tokens starts reading from.

    0, the table's first.
    """
    return 0

KVCache

KVCache(c, window: int | None = None)

Bases: Graph

A layer's keys and values and causal attention over them (attend).

A prompt chunk writes a run of rows; a decode step attends with MHA over one query, which costs the context, not the cache.

Parameters:

Name	Type	Description	Default
window	int | None	The positions before a query it attends to besides its own (None: all). Hugging Face's sliding_window counts the query's own, so pass sliding_window - 1.	None

Raises:

Type	Description
ValueError	A ring with a chunk of other than a power of two.

Source code in iron/lm/blocks.py

def __init__(self, c, window: int | None = None):
    G, D, C = c.num_key_value_heads, c.head_dim, c.prefill_chunk
    self.chunk, self.window, rows = C, window, c.max_seq_len
    self.ring = None if window is None else 1 << (window + C - 1).bit_length()
    if self.ring is not None and 3 * self.ring < rows:
        if C & (C - 1):
            raise ValueError(f"KVCache: a {C}-position chunk is no power of two")
        self.staged = iron.state((self.ring, G, D)), iron.state((self.ring, G, D))
        rows = 2 * self.ring
    else:
        self.ring = None
    self.keys, self.values = iron.state((rows, G, D)), iron.state((rows, G, D))

restart

restart(held: int, ran: int) -> int

Return where a call may start to change the positions from held on.

After ran ran: the chunk boundary before it while the ring still holds the window before that, else the first.

Source code in iron/lm/blocks.py

def restart(self, held: int, ran: int) -> int:
    """Return where a call may start to change the positions from ``held`` on.

    After ``ran`` ran: the chunk boundary before it while the ring
    still holds the window before that, else the first.
    """
    assert self.ring is not None and self.window is not None, "restart needs a ring"
    start = held // self.chunk * self.chunk
    return start if ran - (start - self.window) <= self.ring else 0

attend

attend(call, q, k, v, scale: float | None = None)

Write k and v (rows, groups, head_dim) at the call's positions.

None for a layer reading another's cache. Attends q (rows, heads, head_dim) over them; (rows, heads * head_dim).

Source code in iron/lm/blocks.py

def attend(self, call, q, k, v, scale: float | None = None):
    """Write ``k`` and ``v`` ``(rows, groups, head_dim)`` at the call's positions.

    None for a layer reading another's cache. Attends ``q`` ``(rows,
    heads, head_dim)`` over them; ``(rows, heads * head_dim)``.
    """
    C, (_, G, D), (n, H, _) = self.chunk, self.keys.shape, q.shape
    if self.ring is None:
        for x, cache in ((k, self.keys), (v, self.values)):
            if x is None:
                continue
            if call.prompt:
                cache[call.start : call.start + C][: call.rows] = x
            else:
                cache[call.position] = x.reshape(G, D)
        keys, values = (x[: call.position + 1] for x in (self.keys, self.values))
    else:
        assert self.window is not None, "a ring has a window"
        keys, values, q = self._ring(call, self.ring, self.window, q, k, v)
    o = MHA.step(
        q,
        keys,
        values,
        heads_interleaved=True,
        kv_interleaved=True,
        window=self.window,
        scale=scale,
    )
    if call.prompt and self.ring is not None:
        o = o[: call.rows]
    return o.reshape(n, H * D)

NoCache

NoCache(window: int | None = None, causal: bool = False)

Bases: Graph

No memory: attention over the call's own rows, both ways, or back if causal.

Within window positions either side (None: all), what an encoder's layer holds in a KVCache's place.

Source code in iron/lm/blocks.py

def __init__(self, window: int | None = None, causal: bool = False):
    self.window, self.causal = window, causal

attend

attend(call, q, k, v, scale: float | None = None)

Attend q (rows, heads, head_dim) over k and v.

k and v are (rows, groups, head_dim), the logits scaled by scale (MHA's 1 / sqrt(head_dim) if None); (rows, heads * head_dim).

Source code in iron/lm/blocks.py

def attend(self, call, q, k, v, scale: float | None = None):
    """Attend ``q`` ``(rows, heads, head_dim)`` over ``k`` and ``v``.

    ``k`` and ``v`` are ``(rows, groups, head_dim)``, the logits scaled
    by ``scale`` (MHA's ``1 / sqrt(head_dim)`` if None); ``(rows, heads
    * head_dim)``.
    """
    n, H, D = q.shape
    o = MHA.step(
        q,
        k,
        v,
        heads_interleaved=True,
        kv_interleaved=True,
        causal=self.causal,
        window=self.window,
        scale=scale,
    )
    return o.reshape(n, H * D)

Attention

Attention(c, i: int, norms: str = '', o_bias: bool = False)

Bases: Graph

Self-attention for layer i of a decoder with config c.

q, k and v project, RoPE rotates the first rotary columns when set, and attention runs over the layer's KVCache.

Parameters:

Name	Type	Description	Default
norms	str	Which of q, k and v each head of is RMS-normed before RoPE (q_norm, k_norm; v unweighted).	''

Raises:

Type	Description
ValueError	norms names other than q, k and v.

Source code in iron/lm/blocks.py

def __init__(self, c, i: int, norms: str = "", o_bias: bool = False):
    if set(norms) - set("qkv"):
        raise ValueError(f"Attention: norms {norms!r} names other than q, k and v")
    H, G, D, E = (
        c.num_attention_heads,
        c.num_key_value_heads,
        c.head_dim,
        c.hidden_size,
    )
    self.heads, self.groups, self.head_dim = H, G, D
    self.projections(c)
    self.o_proj = Linear(H * D, E, bias=o_bias, precision=c.precision)
    self.q_norm = RMSNorm(D, c.eps) if "q" in norms else None
    self.k_norm = RMSNorm(D, c.eps) if "k" in norms else None
    self.v_norm, self.eps = "v" in norms, c.eps
    self.scaling: float | None = None
    self.rotary = c.rotary_dim

projections

projections(c) -> None

q_proj, k_proj and v_proj.

Source code in iron/lm/blocks.py

def projections(self, c) -> None:
    """``q_proj``, ``k_proj`` and ``v_proj``."""
    E, D = c.hidden_size, self.head_dim
    linear = partial(Linear, bias=c.attention_bias, precision=c.precision)
    self.q_proj = linear(E, self.heads * D)
    self.k_proj = linear(E, self.groups * D)
    self.v_proj = linear(E, self.groups * D)

project

project(x, **gemv) -> list

Project q, k and v of x, each given the GEMV's gemv.

None for a projection the layer does not hold.

Source code in iron/lm/blocks.py

def project(self, x, **gemv) -> list:
    """Project q, k and v of ``x``, each given the GEMV's ``gemv``.

    None for a projection the layer does not hold.
    """
    return [
        (
            None
            if p is None
            else F.linear(x, p.weight, p.bias, precision=p.precision, **gemv)
        )
        for p in (self.q_proj, self.k_proj, self.v_proj)
    ]

rotate

rotate(x, angles)

Turn x (rows, head_dim) by RoPE.

Only its first rotary columns where given.

Source code in iron/lm/blocks.py

def rotate(self, x, angles):
    """Turn ``x`` ``(rows, head_dim)`` by RoPE.

    Only its first ``rotary`` columns where given.
    """
    if self.rotary is None:
        return F.rope(x, angles)
    x[:, : self.rotary] = F.rope(Copy.step(x[:, : self.rotary]), angles)
    return x

output

output(o: Handle, x: Handle) -> Handle

Return the block's output from the heads o (rows, heads * head_dim).

Of the input x.

Source code in iron/lm/blocks.py

def output(self, o: Handle, x: Handle) -> Handle:
    """Return the block's output from the heads ``o`` ``(rows, heads * head_dim)``.

    Of the input ``x``.
    """
    return self.o_proj(o)

MLP

MLP(c)

Bases: Graph

A feed-forward on c's shape: fc2(act(fc1(x))), act is hidden_act.

Each projection has a bias where the config's mlp_bias says, at precision.

Source code in iron/lm/blocks.py

def __init__(self, c):
    E, M = c.hidden_size, c.intermediate_size
    linear = partial(Linear, bias=c.mlp_bias, precision=c.precision)
    self.fc1 = linear(E, M)
    self.fc2 = linear(M, E)
    self.act = activation(c.hidden_act)

GatedMLP

GatedMLP(c)

Bases: Graph

A gated feed-forward, down_proj(act(gate_proj(x)) * up_proj(x)).

act the config's hidden_act; each projection has a bias when mlp_bias is set.

Source code in iron/lm/blocks.py

def __init__(self, c):
    self.projections(c)
    self.down_proj = Linear(
        c.intermediate_size, c.hidden_size, c.mlp_bias, precision=c.precision
    )
    self.act = activation(c.hidden_act)

projections

projections(c) -> None

gate_proj and up_proj.

Source code in iron/lm/blocks.py

def projections(self, c) -> None:
    """``gate_proj`` and ``up_proj``."""
    E, M = c.hidden_size, c.intermediate_size
    linear = partial(Linear, bias=c.mlp_bias, precision=c.precision)
    self.gate_proj = linear(E, M)
    self.up_proj = linear(E, M)

DecoderLayer

DecoderLayer(c, i: int, self_attn=None, mlp=None, cache=None)

Bases: Graph

Layer i of a pre-norm decoder with config c.

input_layernorm then self_attn over the layer's cache, post_attention_layernorm then mlp, each scaled by residual_multiplier before its residual.

Source code in iron/lm/blocks.py

def __init__(self, c, i: int, self_attn=None, mlp=None, cache=None):
    self.input_layernorm = self.norm(c.hidden_size, c.eps)
    self.self_attn = (
        Attention(c, i, o_bias=c.attention_bias) if self_attn is None else self_attn
    )
    self.post_attention_layernorm = self.norm(c.hidden_size, c.eps)
    self.mlp = GatedMLP(c) if mlp is None else mlp
    self.cache = KVCache(c, c.window) if cache is None else cache
    self.rope: Hashable | None = None
    self.residual_multiplier = 1.0

mix

mix(h: Handle, call) -> Handle

Attend the normed rows h over the layer's cache.

Source code in iron/lm/blocks.py

def mix(self, h: Handle, call) -> Handle:
    """Attend the normed rows ``h`` over the layer's ``cache``."""
    angles = None if self.rope is None else call.ropes[self.rope]
    return self.self_attn(h, self.cache, call, angles)

EncoderLayer

EncoderLayer(c, i: int, self_attn=None)

Bases: DecoderLayer

Layer i of a pre-norm encoder, as in SigLIP and Whisper.

DecoderLayer's two blocks after a LayerNorm, bidirectional attention with no cache (NoCache) and a plain (ungated) MLP.

Source code in iron/lm/blocks.py

def __init__(self, c, i: int, self_attn=None):
    super().__init__(c, i, self_attn, MLP(c), NoCache())

SandwichLayer

SandwichLayer(c, i: int, self_attn=None, mlp=None, cache=None)

Bases: DecoderLayer

DecoderLayer with each block's output normed again before its residual.

(post_attention_layernorm, post_feedforward_layernorm) and mlp's input normed by pre_feedforward_layernorm (Gemma's).

Source code in iron/lm/blocks.py

def __init__(self, c, i: int, self_attn=None, mlp=None, cache=None):
    super().__init__(c, i, self_attn, mlp, cache)
    self.pre_feedforward_layernorm = RMSNorm(c.hidden_size, c.eps)
    self.post_feedforward_layernorm = RMSNorm(c.hidden_size, c.eps)

PleLayer

PleLayer(c, i: int, self_attn=None, mlp=None, cache=None)

Bases: SandwichLayer

SandwichLayer, then its per-layer input ple (rows, ple_dim) gated in.

(per_layer_input_gate, per_layer_projection, post_per_layer_input_norm) and scaled by layer_scalar (Gemma 4's).

Source code in iron/lm/blocks.py

def __init__(self, c, i: int, self_attn=None, mlp=None, cache=None):
    super().__init__(c, i, self_attn, mlp, cache)
    E, P = c.hidden_size, c.ple_dim
    self.per_layer_input_gate = Linear(E, P, False, precision=c.precision)
    self.per_layer_projection = Linear(P, E, False, precision=c.precision)
    self.post_per_layer_input_norm = RMSNorm(E, c.eps)
    self.layer_scalar = iron.param(1)

filterbank

filterbank(fft: int, mels: int, sample_rate: int, max_hz: float) -> ndarray

Compute the (fft // 2 + 1, mels) unnormalized mel triangles, HTK scale.

Over the rfft's bins, from 0 Hz to max_hz, float64, as transformers' mel_filter_bank.

Source code in iron/lm/blocks.py

def filterbank(fft: int, mels: int, sample_rate: int, max_hz: float) -> np.ndarray:
    """Compute the ``(fft // 2 + 1, mels)`` unnormalized mel triangles, HTK scale.

    Over the rfft's bins, from 0 Hz to ``max_hz``, float64, as transformers'
    ``mel_filter_bank``.
    """
    mel = np.linspace(0, 2595 * math.log10(1 + max_hz / 700), mels + 2)
    edges = 700 * (10 ** (mel / 2595) - 1)
    bins = np.linspace(0, sample_rate // 2, fft // 2 + 1)
    slopes = edges[None, :] - bins[:, None]
    widths = np.diff(edges)
    return np.maximum(
        0, np.minimum(-slopes[:, :-2] / widths[:-1], slopes[:, 2:] / widths[1:])
    )

activation

activation(name: str)

Return the function hidden_act name is.

Raises:

Type	Description
ValueError	No operator computes it: only "silu", "gelu" and "gelu_pytorch_tanh" have one.

Source code in iron/lm/blocks.py

def activation(name: str):
    """Return the function ``hidden_act`` ``name`` is.

    Raises:
        ValueError: No operator computes it: only ``"silu"``, ``"gelu"``
            and ``"gelu_pytorch_tanh"`` have one.
    """
    if name == "silu":
        return F.silu
    if name == "gelu":
        return F.gelu
    if name == "gelu_pytorch_tanh":
        return partial(F.gelu, approximate="tanh")
    raise ValueError(
        f"no operator for hidden_act {name!r}; silu, gelu and gelu_pytorch_tanh "
        "have one"
    )

residual_add

residual_add(x, y, residual: float)

x plus y times residual, in x's shape.

Source code in iron/lm/blocks.py

def residual_add(x, y, residual: float):
    """``x`` plus ``y`` times ``residual``, in ``x``'s shape."""
    y = y if y.shape == x.shape else y.reshape(x.shape)
    if residual == 1:
        return x + y
    return AXPY.step(y, x, scalar_factor=residual)

feature_map

Conv2d, a convolution over the rows of a channels-last feature map.

The map is (rows, width * channels).

Conv2d

Conv2d(weight: ndarray, width: int, stride: int = 1, padding: int = 0, border: int = 0, groups: int = 1)

Bases: Graph

A bias-free 2-D convolution of channels-last rows (rows, width * channels).

Any row padding is held in the input. It computes the per-call rows of (out_rows, (out_width + 2 * border) * out_channels) summed in float32. Its output is written to a given y only when it fills whole GEMM blocks.

Parameters:

Name	Type	Description	Default
weight	ndarray	(out_channels, in_channels // groups, kernel, kernel).	required
width	int	The input's pixels a row, its held padding among them.	required
stride	int	The step between output pixels, and rows.	1
padding	int	The zero columns either side of the width; one channel or depthwise only.	0
border	int	The zero pixels either side of each output row; one channel or depthwise only.	0
groups	int	1, or out_channels (depthwise); no other grouping is built.	1

Source code in iron/common/feature_map.py

def __init__(
    self,
    weight: np.ndarray,
    width: int,
    stride: int = 1,
    padding: int = 0,
    border: int = 0,
    groups: int = 1,
):
    w = np.asarray(weight, np.float32)
    C_out, channels_per_group, k, _ = w.shape
    depthwise = groups > 1
    if depthwise and (groups != C_out or channels_per_group != 1):
        raise ValueError(
            "Conv2d: groups > 1 is depthwise only: groups must equal "
            "out_channels, and the weight's own channel axis must be 1"
        )
    C = groups if depthwise else channels_per_group
    banded = C == 1 or depthwise
    if not banded and (padding or border):
        raise ValueError(
            "Conv2d: gathered patches hold no zero padding or border; with "
            "more than one channel (not depthwise) the input holds its "
            "padding columns"
        )
    self.width, self.stride, self.padding, self.border = (
        width,
        stride,
        padding,
        border,
    )
    self.channels, self.out_channels, self.kernel = C, C_out, k
    self.banded, self.depthwise = banded, depthwise
    self.out_width = (width + 2 * padding - k) // stride + 1
    if self.out_width < 1:
        raise ValueError(
            f"Conv2d: a {k}-wide kernel does not fit {width} pixels padded "
            f"by {padding}"
        )
    self.out_line = (self.out_width + 2 * border) * C_out
    if C == 1:
        band = np.zeros((k, width, self.out_width + 2 * border, C_out), np.float32)
        for f in range(self.out_width):
            for b in range(k):
                i = stride * f - padding + b
                if 0 <= i < width:
                    band[:, i, f + border] = w[:, 0, :, b].T
        matrix = band.reshape(k * width, self.out_line)
    elif depthwise:
        band = np.zeros(
            (k, width, C, self.out_width + 2 * border, C_out), np.float32
        )
        channel = np.arange(C)
        for f in range(self.out_width):
            for b in range(k):
                i = stride * f - padding + b
                if 0 <= i < width:
                    band[:, i, channel, f + border, channel] = w[:, 0, :, b].T
        matrix = band.reshape(k * width * C, self.out_line)
    else:
        matrix = w.transpose(0, 2, 3, 1).reshape(C_out, k * k * C)
    self.matrix = iron.param(matrix.astype(bfloat16))

causal

A decoder-only language model on the NPU, as one graph.

A call with no tensor is a decode step; prefill_chunk token ids are a prompt chunk from start, of which a call runs the first rows. Both draw the next token on the device and carry it and position + 1 into the next call.

Call dataclass

Call(*, prompt: bool, start: Any, rows: Any, position: Any, rope_start: Any = None, ropes: dict = dict(), ids: Any = None, token: Any = None, embeddings: Any = None)

One call's values, built once in body and read by every memory module.

A memory module is one keeping state from call to call (KVCache, Checkpoints, Rope): a prompt chunk's rows from start, or the decode step's one at position.

Checkpoints dataclass

Checkpoints(slots: State, chunk: int)

A state a decoder carries from position to position, a slot per prompt-chunk boundary.

Slot k holds it after the positions up to k * chunk run so far, slot 0 (never written) before the sequence.

of classmethod

of(shape: tuple[int, ...], max_seq_len: int, chunk: int, dtype=bfloat16) -> Checkpoints

Return a state of shape.

A slot for every boundary of chunk-position prompt chunks up to max_seq_len positions.

Raises:

Type	Description
ValueError	chunk is not a power of two, which a decode step's slot (position // chunk) needs.

Source code in iron/lm/causal.py

@classmethod
def of(
    cls, shape: tuple[int, ...], max_seq_len: int, chunk: int, dtype=bfloat16
) -> "Checkpoints":
    """Return a state of ``shape``.

    A slot for every boundary of ``chunk``-position prompt chunks up to
    ``max_seq_len`` positions.

    Raises:
        ValueError: ``chunk`` is not a power of two, which a decode step's
            slot (``position // chunk``) needs.
    """
    if chunk < 1 or chunk & (chunk - 1):
        raise ValueError(
            f"Checkpoints: a {chunk}-position chunk is no power of two"
        )
    slots = (max_seq_len - 1) // chunk + 2
    return cls(iron.state((slots, *shape), dtype), chunk)

at

at(call: Call) -> tuple[Handle, Handle]

Return the slot call reads and the slot it writes, in a graph's body.

The one read holds the state before its first position, the one written the state after its last. A call ends within the chunk it starts in (CausalLM splits them).

Source code in iron/lm/causal.py

def at(self, call: Call) -> tuple[Handle, Handle]:
    """Return the slot ``call`` reads and the slot it writes, in a graph's body.

    The one read holds the state before its first position, the one
    written the state after its last. A call ends within the chunk it
    starts in (``CausalLM`` splits them).
    """
    C = self.chunk
    return self.slots[(call.start - 1) // C + 1], self.slots[call.position // C + 1]

restart

restart(held: int, ran: int) -> int

Where a call may start to change the positions from held on, after ran positions ran.

There where the last call ended, else at the boundary before it.

Source code in iron/lm/causal.py

def restart(self, held: int, ran: int) -> int:
    """Where a call may start to change the positions from ``held`` on, after ``ran`` positions ran.

    There where the last call ended, else at the boundary before it.
    """
    return held if held == ran else held // self.chunk * self.chunk

slot

slot(position: int) -> int

Return the slot holding the state after the first position positions.

Source code in iron/lm/causal.py

def slot(self, position: int) -> int:
    """Return the slot holding the state after the first ``position`` positions."""
    return 0 if position == 0 else (position - 1) // self.chunk + 1

CausalConv1d

CausalConv1d(c, channels: int, kernel: int, bias: bool = False)

Bases: Graph

A depth-wise causal convolution of kernel taps over a call's rows.

Each of channels its own taps (weight, torch's (channels, 1, kernel)), plus bias where it holds one. The kernel - 1 rows before a call are its state, a slot per prompt-chunk boundary; CausalLM gives every one of a model the scratch F.causal_conv1d stages a call in (window, tail).

Raises:

Type	Description
ValueError	prefill_chunk is smaller than kernel - 1, the history the staging window must hold.

Source code in iron/lm/causal.py

def __init__(self, c, channels: int, kernel: int, bias: bool = False):
    if c.prefill_chunk < kernel - 1:
        raise ValueError(
            f"CausalConv1d: prefill_chunk {c.prefill_chunk} is smaller than "
            f"kernel - 1 ({kernel - 1})"
        )
    self.weight = iron.param(channels, 1, kernel)
    self.bias = iron.param(channels) if bias else None
    self.state = Checkpoints.of(
        (kernel - 1, channels), c.max_seq_len, c.prefill_chunk
    )

Decoder

Decoder(c: DecoderConfig, layer: type[DecoderLayer])

Bases: Graph

A decoder's modules on c's shape, as a checkpoint's model names them.

embed_tokens, layers, layer i a layer(c, i), and the final norm.

Source code in iron/lm/causal.py

def __init__(self, c: DecoderConfig, layer: type[DecoderLayer]):
    self.embed_tokens = Embedding(c.vocab_size, c.hidden_size)
    self.layers = [layer(c, i) for i in range(c.num_hidden_layers)]
    self.norm = RMSNorm(c.hidden_size, c.eps)

CausalLM

CausalLM(config: DecoderConfig, weights)

Bases: Model

A decoder on config's shape: model, lm_head and rotary_emb.

model is a decoder_class of layer_class layers each holding its own KVCache; rotary_emb is the RoPE table, None with no rope_theta.

Source code in iron/lm/causal.py

def __init__(self, config: DecoderConfig, weights):
    self.config = config
    self.weights = weights
    self.rotary_emb = None if config.rope_theta is None else self.rotary(config)
    # Zero until the host writes draws: a row of zeros is greedy.
    self.draws = iron.state((config.max_seq_len, ROW_WORDS), np.int32)
    self.drawn = iron.state((config.max_seq_len,), np.int32)
    # The versions ``use`` runs, which every session calls.
    self.prompt_chunk: CompiledGraph | None = None
    self.decode_step: CompiledGraph | None = None
    self.placed_chunk: CompiledGraph | None = None
    self.placeholders: tuple[int, ...] = ()
    # By the version a device loop starts from, its loop (``loop``).
    self._loops: dict[int, CarriedLoop] = {}
    # The session whose tokens the caches hold, and the one ``logits``
    # and ``generate`` run in.
    self.holder: Session | None = None
    self.session = Session(self)
    self.model = self.decoder_class(config, self.layer_class)
    for layer in self.model.layers:
        layer.rope = self.rotary_emb
    self.lm_head = Linear(
        config.hidden_size, config.vocab_size, False, precision=config.precision
    )
    if config.tie_word_embeddings:
        self.lm_head.weight = self.model.embed_tokens.weight
    self.load_state_dict(renamed(weights, self.renames))
    if config.vocab_size % 2:
        # An odd vocabulary's logits end mid-word; a zero row makes them
        # whole words, and Sample draws none from it.
        head = np.asarray(self.lm_head.weight)
        self.lm_head.weight = iron.param(np.pad(head, ((0, 1), (0, 0))))
    self.layers = self.model.layers
    # The scratch every causal conv of a shape stages its calls in.
    scratch = {}
    for conv in {id(p): p for _, _, p, _ in walk(self)}.values():
        if isinstance(conv, CausalConv1d):
            (K, C), chunk = conv.taps.shape, config.prefill_chunk
            if (K, C) not in scratch:
                scratch[K, C] = iron.state((2, chunk, C)), iron.state((K - 1, C))
            conv.window, conv.tail = scratch[K, C]

recurrent property

recurrent: list[Checkpoints | KVCache]

Each state a layer carries from position to position, which restarts only at some (restart).

Checkpoints, and a KVCache keeping a ring.

device_loop property

device_loop: bool

Whether generate can draw on the device: a full-ELF decode step.

draws_at_once property

draws_at_once: bool

A device loop's: generate reads its draws once the last is drawn.

layer

layer(call: Call, i: int, w, x)

Layer i, the module w, on x.

Source code in iron/lm/causal.py

def layer(self, call: Call, i: int, w, x):
    """Layer ``i``, the module ``w``, on ``x``."""
    return w(x, call)

head

head(x)

Run the final norm, then the output head (the embedding, if tied).

Source code in iron/lm/causal.py

def head(self, x):
    """Run the final norm, then the output head (the embedding, if tied)."""
    return self.lm_head(self.model.norm(x))

positions

positions(x, call: Call)

x, the rows call runs, with the model's absolute positions placed.

Unchanged by default.

Source code in iron/lm/causal.py

def positions(self, x, call: Call):
    """``x``, the rows ``call`` runs, with the model's absolute positions placed.

    Unchanged by default.
    """
    return x

lookup

lookup(ids)

Return the embedding rows of ids.

Bounded as ids is, or of the decode step's token.

Source code in iron/lm/causal.py

def lookup(self, ids):
    """Return the embedding rows of ``ids``.

    Bounded as ``ids`` is, or of the decode step's token.
    """
    return self.model.embed_tokens(ids)

from_rows

from_rows(x, ids, *, token, position, start, rows, rope_start)

Run a prompt chunk's call on its embedding rows x, bounded to rows.

In place of the lookup of its ids; token is the decode step's, unread here.

Source code in iron/lm/causal.py

def from_rows(self, x, ids, *, token, position, start, rows, rope_start):
    """Run a prompt chunk's call on its embedding rows ``x``, bounded to ``rows``.

    In place of the lookup of its ``ids``; ``token`` is the decode
    step's, unread here.
    """
    return self.forward(
        x,
        Call(
            prompt=True,
            start=start,
            rows=rows,
            position=position,
            rope_start=rope_start,
            ids=ids,
        ),
    )

forward

forward(x: Handle, call: Call)

Run the layers, the head and the draw over x, the rows call runs.

Each memory module reads call.

Source code in iron/lm/causal.py

def forward(self, x: Handle, call: Call):
    """Run the layers, the head and the draw over ``x``, the rows ``call`` runs.

    Each memory module reads ``call``.
    """
    c = self.config
    call.embeddings = x
    for rope in (layer.rope for layer in self.layers):
        if rope is not None and rope not in call.ropes:
            call.ropes[rope] = rope.rows(call)
    x = self.positions(x, call)
    for i, weights in enumerate(self.layers):
        x = self.layer(call, i, weights, x)
    if call.prompt:
        x = Copy.step(x[call.rows - 1]).reshape(
            1, c.hidden_size
        )  # the last row's logits alone
    return self.draw(self.head(x), call)

draw

draw(logits, call: Call)

Return what a call returns: logits and the token drawn from them.

logits is head's output; the token is drawn from it on the device, carried with the next position.

Source code in iron/lm/causal.py

def draw(self, logits, call: Call):
    """Return what a call returns: ``logits`` and the token drawn from them.

    ``logits`` is ``head``'s output; the token is drawn from it on the
    device, carried with the next position.
    """
    c = self.config
    _, drawn = Sample.step(
        logits.reshape(c.vocab_size + c.vocab_size % 2),
        self.draws,
        self.drawn,
        row=call.position * ROW_WORDS,
        at=call.position,
        drawable=c.vocab_size,
    )
    return logits, iron.carry(token=drawn, position=call.position + 1)

rotary

rotary(config: DecoderConfig) -> Rope

rotary_emb, the RoPE table every layer holds: Rope's.

Source code in iron/lm/causal.py

def rotary(self, config: DecoderConfig) -> Rope:
    """``rotary_emb``, the RoPE table every layer holds: ``Rope``'s."""
    return Rope(config)

rope_row

rope_row(n: int, begin: int) -> int

Return the RoPE table row a chunk from begin reads first.

Of an n-token prompt (Rope.table_start).

Source code in iron/lm/causal.py

def rope_row(self, n: int, begin: int) -> int:
    """Return the RoPE table row a chunk from ``begin`` reads first.

    Of an ``n``-token prompt (``Rope.table_start``).
    """
    rope = self.rotary_emb
    return begin + (0 if rope is None else rope.table_start(n))

tune_arguments classmethod

tune_arguments(parser) -> None

Add the flags tune_calls follows from (tune.arguments).

Source code in iron/lm/causal.py

@classmethod
def tune_arguments(cls, parser) -> None:
    """Add the flags ``tune_calls`` follows from (``tune.arguments``)."""
    from .tune import arguments  # which imports this module

    arguments(parser)

tune_calls classmethod

tune_calls(parser, args) -> list

tune.calls on args, built on weights it never reads.

Source code in iron/lm/causal.py

@classmethod
def tune_calls(cls, parser, args) -> list:
    """``tune.calls`` on ``args``, built on weights it never reads."""
    from .tune import calls  # which imports this module

    weights = weights_file(args.checkpoint)
    config = cls.default_config(weights)
    assert isinstance(config, DecoderConfig), "a decoder's"
    model = cls(config, Checkpoint(weights).tree("unread"))
    L, C = config.max_seq_len, config.prefill_chunk
    if not C <= args.context <= L or args.context % C:
        parser.error(
            f"--context {args.context} is not whole {C}-token chunks within "
            f"max_seq_len {L}"
        )
    sampler = Sampler(args.temperature, args.top_k)
    return calls(
        model, args.position, args.token, args.context, args.points, sampler
    )

shapes

shapes() -> list[dict]

Return the input shapes of each version.

The decode step's (none), then the prompt chunk's token ids.

Source code in iron/lm/causal.py

def shapes(self) -> list[dict]:
    """Return the input shapes of each version.

    The decode step's (none), then the prompt chunk's token ids.
    """
    return [{}, dict(ids=((self.config.prefill_chunk,), np.int32))]

loaded

loaded(versions: list) -> None

Run logits and generate on the decode step and the prompt chunk.

See use.

Source code in iron/lm/causal.py

def loaded(self, versions: list) -> None:
    """Run ``logits`` and ``generate`` on the decode step and the prompt chunk.

    See ``use``.
    """
    self.use(*versions)

use

use(decode: CompiledGraph, prompt: CompiledGraph, placed: CompiledGraph | None = None, placeholders: tuple[int, ...] = ()) -> None

Run logits and generate on these versions, in one arena.

The versions are this model's own, or those of a graph recording it in their bodies (a multimodal model's); each prompt version feeds decode.

Parameters:

Name	Type	Description	Default
decode	CompiledGraph	The decode step.	required
prompt	CompiledGraph	A prompt chunk, on ids.	required
placed	CompiledGraph | None	A prompt chunk holding any of placeholders, on its ids and before, the count of each placeholder before the chunk.	None
placeholders	tuple[int, ...]	The placeholder tokens, in before's order.	()

Source code in iron/lm/causal.py

def use(
    self,
    decode: CompiledGraph,
    prompt: CompiledGraph,
    placed: CompiledGraph | None = None,
    placeholders: tuple[int, ...] = (),
) -> None:
    """Run ``logits`` and ``generate`` on these versions, in one arena.

    The versions are this model's own, or those of a graph recording it
    in their bodies (a multimodal model's); each prompt version feeds
    ``decode``.

    Args:
        decode: The decode step.
        prompt: A prompt chunk, on ``ids``.
        placed: A prompt chunk holding any of ``placeholders``, on its
            ``ids`` and ``before``, the count of each placeholder before
            the chunk.
        placeholders: The placeholder tokens, in ``before``'s order.
    """
    self.decode_step, self.prompt_chunk = decode, prompt
    self.placed_chunk, self.placeholders = placed, tuple(placeholders)
    self._loops, self.holder = {}, None
    if not decode.plan.full_elf:
        logger.info("%s", decode.plan.report("decode"))

call_values

call_values() -> dict

Return the per-call values every call gives beyond the standard ones.

The standard ones are the token, position, start, rows and rope_start; none beyond them by default.

Source code in iron/lm/causal.py

def call_values(self) -> dict:
    """Return the per-call values every call gives beyond the standard ones.

    The standard ones are the token, position, start, rows and
    rope_start; none beyond them by default.
    """
    return {}

release

release() -> None

Free the device loops' runs, then the graph's memory.

Source code in iron/lm/causal.py

def release(self) -> None:
    """Free the device loops' runs, then the graph's memory."""
    self._loops.clear()
    super().release()

loop

loop(version: CompiledGraph) -> CarriedLoop

Return the device loop of decode steps version starts, made once.

Its runs are bound to the versions, which every session shares.

Source code in iron/lm/causal.py

def loop(self, version: CompiledGraph) -> CarriedLoop:
    """Return the device loop of decode steps ``version`` starts, made once.

    Its runs are bound to the versions, which every session shares.
    """
    if id(version) not in self._loops:
        assert self.decode_step is not None, "loaded with the version it starts"
        self._loops[id(version)] = CarriedLoop(version, self.decode_step, depth=2)
    return self._loops[id(version)]

logits

logits(tokens) -> ndarray

Return the logits after the last of tokens, the whole history.

Run in session (Session.logits).

Source code in iron/lm/causal.py

def logits(self, tokens) -> np.ndarray:
    """Return the logits after the last of ``tokens``, the whole history.

    Run in ``session`` (``Session.logits``).
    """
    return self.session.logits(tokens)

forget

forget(position: int) -> None

Hold nothing the caches have from position on.

Whichever session's tokens they hold (Session.forget).

Source code in iron/lm/causal.py

def forget(self, position: int) -> None:
    """Hold nothing the caches have from ``position`` on.

    Whichever session's tokens they hold (``Session.forget``).
    """
    if self.holder is not None:
        self.holder.forget(position)

encoder

An encoder on the NPU as one graph: a sequence of at most max_tokens rows in.

Attended bidirectionally, compiled once for every length.

Encoder looks up token ids and runs forward over them; a Tower is an encoder whose rows come from a front end and whose output is soft tokens in a text model's space.

Encoder

Encoder(config: Any, max_tokens: int)

Bases: Model

An encoder over max_tokens rows at most, one compile for every length.

layers, a layer_class(config, i) each, which a subclass loads from its checkpoint; it gives lookup and head.

Parameters:

Name	Type	Description	Default
config	Any	The model's shape: num_hidden_layers layers, and inputs reads its vocab_size.	required

Source code in iron/lm/encoder.py

def __init__(self, config: Any, max_tokens: int):
    self.config = config
    self.layers = [
        self.layer_class(config, i) for i in range(config.num_hidden_layers)
    ]
    self.max_tokens = max_tokens
    # The rows a call gives, the first ``soft(values)`` of them valid.
    self.soft_rows = max_tokens
    # The norm module after the last layer.
    self.final_norm = None

load_layers

load_layers(tree, layers: str, norm: str = '') -> None

Load layer i from tree, renamed by renames, at {layers}{i}..

Then loads final_norm, where there is one, at norm.

Source code in iron/lm/encoder.py

def load_layers(self, tree, layers: str, norm: str = "") -> None:
    """Load layer ``i`` from ``tree``, renamed by ``renames``, at ``{layers}{i}.``.

    Then loads ``final_norm``, where there is one, at ``norm``.
    """
    tree = renamed(tree, self.renames)
    for i, layer in enumerate(self.layers):
        layer.load_state_dict(tree, path=f"{layers}{i}.")
    if self.final_norm is not None:
        self.final_norm.load_state_dict(tree, path=norm)

lookup

lookup(ids)

Return the embedding rows of ids, bounded as ids is.

Source code in iron/lm/encoder.py

def lookup(self, ids):
    """Return the embedding rows of ``ids``, bounded as ``ids`` is."""
    raise NotImplementedError(f"{type(self).__name__} defines no lookup()")

angles

angles(rows) -> dict | None

Return the RoPE angles of the first rows positions, by the rope each layer names.

Call.ropes; None: no RoPE.

Source code in iron/lm/encoder.py

def angles(self, rows, /) -> dict | None:
    """Return the RoPE angles of the first ``rows`` positions, by the ``rope`` each layer names.

    ``Call.ropes``; None: no RoPE.
    """
    return None

positions

positions(x, call: Call)

Return x, the rows call runs, with the model's absolute positions placed.

Unchanged by default.

Source code in iron/lm/encoder.py

def positions(self, x, call: Call):
    """Return ``x``, the rows ``call`` runs, with the model's absolute positions placed.

    Unchanged by default.
    """
    return x

layer

layer(call: Call, i: int, weights, x)

Layer i, the module weights, on x.

Source code in iron/lm/encoder.py

def layer(self, call: Call, i: int, weights, x, /):
    """Layer ``i``, the module ``weights``, on ``x``."""
    return weights(x, call)

pool

pool(x, rows)

Return the rows head reads from the rows valid rows of x.

All of them by default.

Source code in iron/lm/encoder.py

def pool(self, x, rows, /):
    """Return the rows ``head`` reads from the ``rows`` valid rows of ``x``.

    All of them by default.
    """
    return x

forward

forward(x: Handle, rows, ropes, ids=None)

Run the layers, the final norm, pool and head over the rows x.

Of ids (None for a tower's), bounded to their first rows, rotated by ropes (Call.ropes).

Source code in iron/lm/encoder.py

def forward(self, x: Handle, rows, ropes, ids=None):
    """Run the layers, the final norm, ``pool`` and ``head`` over the rows ``x``.

    Of ``ids`` (None for a tower's), bounded to their first ``rows``,
    rotated by ``ropes`` (``Call.ropes``).
    """
    call = Call(
        prompt=True,
        start=0,
        rows=rows,
        position=None,
        ropes=ropes,
        ids=ids,
        embeddings=x,
    )
    x = self.positions(x, call)
    for i, weights in enumerate(self.layers):
        x = self.layer(call, i, weights, x)
    if self.final_norm is not None:
        x = self.final_norm(x)
    return self.head(self.pool(x, rows))

from_rows

from_rows(x, ids, *, rows)

Run a call on its embedding rows x, bounded to rows.

In place of the lookup of its ids.

Source code in iron/lm/encoder.py

def from_rows(self, x, ids, *, rows):
    """Run a call on its embedding rows ``x``, bounded to ``rows``.

    In place of the lookup of its ``ids``.
    """
    return self.forward(x, rows, self.angles(rows), ids)

shapes

shapes() -> list[dict]

Return the input shapes of each version: one, max_tokens rows.

Source code in iron/lm/encoder.py

def shapes(self) -> list[dict]:
    """Return the input shapes of each version: one, ``max_tokens`` rows."""
    return [dict(ids=((self.max_tokens,), np.int32))]

longest

longest() -> dict

Return the per-call values of the longest call: every row real.

Source code in iron/lm/encoder.py

def longest(self) -> dict:
    """Return the per-call values of the longest call: every row real."""
    return dict(rows=self.max_tokens)

soft

soft(values: dict) -> int

Return the valid rows a call with per-call values gives.

Its rows by default.

Source code in iron/lm/encoder.py

def soft(self, values: dict) -> int:
    """Return the valid rows a call with per-call ``values`` gives.

    Its ``rows`` by default.
    """
    return values["rows"]

inputs

inputs(tokens) -> tuple[ndarray, dict]

Return a call's ids, tokens padded to max_tokens, and its per-call values.

Raises:

Type	Description
ValueError	tokens is empty or longer than max_tokens.
IndexError	A token is outside [0, vocab_size).

Source code in iron/lm/encoder.py

def inputs(self, tokens) -> tuple[np.ndarray, dict]:
    """Return a call's ``ids``, ``tokens`` padded to ``max_tokens``, and its per-call values.

    Raises:
        ValueError: ``tokens`` is empty or longer than ``max_tokens``.
        IndexError: A token is outside ``[0, vocab_size)``.
    """
    tokens = np.asarray(tokens)
    n, V = len(tokens), self.config.vocab_size
    if not 0 < n <= self.max_tokens:
        raise ValueError(f"{n} tokens do not fit {self.max_tokens} rows")
    outside = tokens[(tokens < 0) | (tokens >= V)]
    if outside.size:
        raise IndexError(f"tokens {outside} are outside [0, {V})")
    ids = np.zeros(self.max_tokens, np.int32)
    ids[:n] = tokens
    return ids, dict(rows=n)

Tower

Tower(config, max_tokens: int, soft_rows: int)

Bases: Encoder

An encoder whose rows come from a front end rather than token ids.

Its output is soft tokens in a text model's space, text_dim wide, which take a placeholder's places.

Parameters:

Name	Type	Description	Default
soft_rows	int	The rows of soft tokens a call returns, the first soft of them valid.	required

Source code in iron/lm/encoder.py

def __init__(self, config, max_tokens: int, soft_rows: int):
    super().__init__(config, max_tokens)
    self.soft_rows = soft_rows

inputs

inputs(*raw: Any, **kwargs: Any) -> tuple[ndarray, dict]

Return a call's input buffer and per-call values for raw.

Source code in iron/lm/encoder.py

def inputs(self, *raw: Any, **kwargs: Any) -> tuple[np.ndarray, dict]:
    """Return a call's input buffer and per-call values for ``raw``."""
    raise NotImplementedError(f"{type(self).__name__} defines no inputs()")

calls

calls(*raw) -> list[tuple[ndarray, dict]]

Return each call raw takes, in turn, its input buffer and per-call values.

Each call's soft tokens after the last's: one, inputs's, by default.

Source code in iron/lm/encoder.py

def calls(self, *raw) -> list[tuple[np.ndarray, dict]]:
    """Return each call ``raw`` takes, in turn, its input buffer and per-call values.

    Each call's soft tokens after the last's: one, ``inputs``'s, by
    default.
    """
    return [self.inputs(*raw)]

soft_tokens

soft_tokens(*raw) -> ndarray

Return the soft tokens of raw, as calls takes it, (soft, text_dim) float32.

Source code in iron/lm/encoder.py

def soft_tokens(self, *raw) -> np.ndarray:
    """Return the soft tokens of ``raw``, as ``calls`` takes it, ``(soft, text_dim)`` float32."""
    calls = self.calls(*raw)
    for buffer, values in calls:
        out = self(buffer, **values)
    out = out.numpy().reshape(-1, self.config.text_dim)
    return np.asarray(out[: sum(self.soft(v) for _, v in calls)], np.float32)

multimodal

A text model whose placeholder tokens take its towers' soft tokens.

As Hugging Face's masked_scatter places them: the j-th of a token's placeholders takes its tower's j-th soft token.

Multimodal

Multimodal(text: Encoder | CausalLM, towers: dict[int, Tower])

Bases: Model

A text model and its towers, keyed by placeholder token, as one graph.

A tower writes its soft tokens into merged, and placed scatters them into the text model's rows by placeholder position.

Parameters:

Name	Type	Description	Default
towers	dict[int, Tower]	Each tower, by the token its placeholders are.	required

Raises:

Type	Description
TypeError	Two inputs or values share a name, or a tower takes other than one input.

Source code in iron/lm/multimodal.py

def __init__(self, text: Encoder | CausalLM, towers: dict[int, Tower]):
    self.text = text
    self.towers = towers
    P = inspect.Parameter
    parameters = list(text.signature().parameters.values())
    self.text_inputs = [p.name for p in parameters if p.kind is not P.KEYWORD_ONLY]
    self.text_values = [p.name for p in parameters if p.kind is P.KEYWORD_ONLY]
    self.tower_inputs, self.tower_values = {}, {}
    for token, tower in towers.items():
        parameters = list(tower.signature().parameters.values())
        inputs = [p.name for p in parameters if p.kind is not P.KEYWORD_ONLY]
        if len(inputs) != 1:
            raise TypeError(f"Multimodal: a tower takes one input, not {inputs}")
        self.tower_inputs[token] = inputs[0]
        self.tower_values[token] = [
            p.name for p in parameters if p.kind is P.KEYWORD_ONLY
        ]
    # Every tower's per-call values, zero: those of a call none runs.
    self.idle = {
        f"{self.tower_inputs[token]}_{v}": 0
        for token in towers
        for v in self.tower_values[token]
    }
    (ids,) = [s["ids"] for s in text.shapes() if "ids" in s]
    (rows,), _ = ids
    self.at = {}
    for token, tower in towers.items():
        self.at[token] = rows
        rows += tower.soft_rows
    self.merged = iron.state((rows, text.config.hidden_size))
    # What each tower ran on last, so a causal model's caches outlive a
    # call with the same.
    self._last_raw: dict[int, tuple] = {}
    self.signature()

shapes

shapes() -> list[dict]

Return the input shapes of each version.

The text model's, first so its decode step leads (Model.shapes), each tower's, and the ids with their placeholders (and a causal model's counts).

Source code in iron/lm/multimodal.py

def shapes(self) -> list[dict]:
    """Return the input shapes of each version.

    The text model's, first so its decode step leads (``Model.shapes``),
    each tower's, and the ids with their placeholders (and a causal
    model's counts).
    """
    text = self.text.shapes()
    (ids,) = [s["ids"] for s in text if "ids" in s]
    towers = [s for t in self.towers.values() for s in t.shapes()]
    placed = dict(placed=ids)
    if isinstance(self.text, CausalLM):
        placed["before"] = ((len(self.towers),), np.int32)
    return [*text, *towers, placed]

release

release() -> None

Free the text model's loops, then every version's and the arena's memory.

Source code in iron/lm/multimodal.py

def release(self) -> None:
    """Free the text model's loops, then every version's and the arena's memory."""
    self.text.release()
    super().release()

loaded

loaded(versions: list) -> None

Hand a causal text model its versions (use).

Source code in iron/lm/multimodal.py

def loaded(self, versions: list) -> None:
    """Hand a causal text model its versions (``use``)."""
    if isinstance(self.text, CausalLM):
        decode_step, prompt_chunk, *_, placed = versions
        self.text.use(decode_step, prompt_chunk, placed, tuple(self.at))

longest

longest() -> dict

Return the per-call values of the longest call.

The text model's and each tower's.

Source code in iron/lm/multimodal.py

def longest(self) -> dict:
    """Return the per-call values of the longest call.

    The text model's and each tower's.
    """
    assert isinstance(self.text, Encoder), "an encoder text model's"
    values = dict(self.text.longest())
    for token, tower in self.towers.items():
        name = self.tower_inputs[token]
        values |= {f"{name}_{v}": x for v, x in tower.longest().items()}
    return values

tower_calls

tower_calls(tokens: Sequence[int] | ndarray, raw: dict[int, tuple]) -> list[tuple[tuple, dict]]

Each tower version raw runs, in turn.

Its inputs and every tower's per-call values, zero for a tower not running.

Parameters:

Name	Type	Description	Default
tokens	Sequence[int] | ndarray	The token ids, each placeholder run as long as its tower's soft tokens.	required
raw	dict[int, tuple]	Each tower's raw input, as its calls takes it, by placeholder token.	required

Raises:

Type	Description
ValueError	A placeholder's count is not its tower's soft tokens.

Source code in iron/lm/multimodal.py

def tower_calls(
    self, tokens: Sequence[int] | np.ndarray, raw: dict[int, tuple]
) -> list[tuple[tuple, dict]]:
    """Each tower version ``raw`` runs, in turn.

    Its inputs and every tower's per-call values, zero for a tower not
    running.

    Args:
        tokens: The token ids, each placeholder run as long as its
            tower's soft tokens.
        raw: Each tower's raw input, as its ``calls`` takes it, by
            placeholder token.

    Raises:
        ValueError: A placeholder's count is not its tower's soft tokens.
    """
    tokens = np.asarray(tokens)
    names = list(self._parameters.inputs)
    calls = []
    for token, tower in self.towers.items():
        name, soft = self.tower_inputs[token], 0
        for buffer, own in tower.calls(*raw[token]) if token in raw else []:
            soft += tower.soft(own)
            inputs = tuple(buffer if n == name else None for n in names)
            calls.append(
                (inputs, self.idle | {f"{name}_{v}": x for v, x in own.items()})
            )
        places = np.count_nonzero(tokens == token)
        if places != soft:
            raise ValueError(
                f"{places} placeholders {token} for {soft} soft tokens"
            )
    return calls

calls

calls(tokens, raw: dict[int, tuple]) -> list[tuple[tuple, dict]]

Build an encoder text model's call: each version it runs, in turn.

Its inputs and per-call values (tower_calls).

Raises:

Type	Description
TypeError	The text model is causal: logits, generate.

Source code in iron/lm/multimodal.py

def calls(self, tokens, raw: dict[int, tuple]) -> list[tuple[tuple, dict]]:
    """Build an encoder text model's call: each version it runs, in turn.

    Its inputs and per-call values (``tower_calls``).

    Raises:
        TypeError: The text model is causal: ``logits``, ``generate``.
    """
    if not isinstance(self.text, Encoder):
        raise TypeError(f"{type(self.text).__name__} runs logits() or generate()")
    ids, text = self.text.inputs(tokens)
    towers = [
        (inputs, values | text) for inputs, values in self.tower_calls(tokens, raw)
    ]
    names = list(self._parameters.inputs)
    last = "placed" if towers else "ids"
    return [
        *towers,
        (tuple(ids if n == last else None for n in names), text | self.idle),
    ]

run

run(tokens, raw: dict[int, tuple])

Return an encoder text model's output for tokens.

raw at their placeholders (as calls takes them), each version run in turn.

Source code in iron/lm/multimodal.py

def run(self, tokens, raw: dict[int, tuple]):
    """Return an encoder text model's output for ``tokens``.

    ``raw`` at their placeholders (as ``calls`` takes them), each
    version run in turn.
    """
    for inputs, values in self.calls(tokens, raw):
        out = self(*inputs, **values)
    return out

place

place(tokens, raw: dict[int, tuple]) -> None

Run each tower raw names.

So a causal text model's prompt chunks find its soft tokens in merged, unless each ran on the same last time, else forgetting its caches from the first stale placeholder.

Source code in iron/lm/multimodal.py

def place(self, tokens, raw: dict[int, tuple]) -> None:
    """Run each tower ``raw`` names.

    So a causal text model's prompt chunks find its soft tokens in
    ``merged``, unless each ran on the same last time, else forgetting
    its caches from the first stale placeholder.
    """
    tokens = np.asarray(tokens)
    calls = self.tower_calls(tokens, raw)
    same = raw.keys() == self._last_raw.keys() and all(
        len(raw[t]) == len(self._last_raw[t])
        and all(np.array_equal(a, b) for a, b in zip(raw[t], self._last_raw[t]))
        for t in raw
    )
    if same:
        return
    text = {v: 0 for v in self.text_values}
    for inputs, values in calls:
        self(*inputs, **values, **text)
    self._last_raw = {t: tuple(np.array(a) for a in raw[t]) for t in raw}
    places = np.flatnonzero(np.isin(tokens, list(self.towers)))
    assert isinstance(self.text, CausalLM), "a causal text model's"
    self.text.forget(int(places[0]) if places.size else tokens.size)

logits

logits(tokens, raw: dict[int, tuple]) -> ndarray

Return a causal text model's logits after tokens.

raw at their placeholders (CausalLM.logits).

Source code in iron/lm/multimodal.py

def logits(self, tokens, raw: dict[int, tuple]) -> np.ndarray:
    """Return a causal text model's logits after ``tokens``.

    ``raw`` at their placeholders (``CausalLM.logits``).
    """
    self.place(tokens, raw)
    assert isinstance(self.text, CausalLM), "a causal text model's"
    return self.text.logits(tokens)

model

Model, a model's graph, every version compiled and loaded before the first call.

Built from a checkpoint's files (load), which generates from a prompt; OnHost is the model with no device; and Tokenizer is a prompt's text to its ids.

Tokenizer

Tokenizer(encoding: TokenizerJson, chat_template: str | None = None, bos_token: str = '')

Text to a prompt's token ids, as transformers' tokenizer begins them, and back.

Over encoding, a TokenizerJson. chat_template is the checkpoint's own.

Source code in iron/lm/model.py

def __init__(
    self,
    encoding: "TokenizerJson",
    chat_template: str | None = None,
    bos_token: str = "",
):
    self.encoding = encoding
    self.chat_template, self.bos_token = chat_template, bos_token

chat

chat(messages: list[dict]) -> list[int]

Return the ids of messages, each {"role": ..., "content": ...}.

As chat_template renders them (jinja2, the models extra), the assistant's turn opened after them.

Raises:

Type	Description
ValueError	The checkpoint has no chat template.
ImportError	jinja2 is not installed.
TemplateError	The template refuses messages (its raise_exception).

Source code in iron/lm/model.py

def chat(self, messages: list[dict]) -> list[int]:
    """Return the ids of ``messages``, each ``{"role": ..., "content": ...}``.

    As ``chat_template`` renders them (jinja2, the ``models`` extra), the
    assistant's turn opened after them.

    Raises:
        ValueError: The checkpoint has no chat template.
        ImportError: jinja2 is not installed.
        jinja2.TemplateError: The template refuses ``messages``
            (its ``raise_exception``).
    """
    if self.chat_template is None:
        raise ValueError("the checkpoint has no chat template")
    _models_extra("jinja2")
    from jinja2 import TemplateError  # the models extra's: optional

    def raise_exception(message):
        raise TemplateError(message)

    text = _chat_template(self.chat_template).render(
        messages=messages,
        add_generation_prompt=True,
        bos_token=self.bos_token,
        raise_exception=raise_exception,
        strftime_now=lambda f: datetime.datetime.now().strftime(f),
    )
    return self.encoding.encode(text)

TokenizerJson

TokenizerJson(path)

A Hugging Face tokenizer.json, over tokenizers (not transformers).

Text to token ids, and back.

Source code in iron/lm/model.py

def __init__(self, path):
    tokenizers = _models_extra("tokenizers")
    self._tokenizer = tokenizers.Tokenizer.from_file(str(path))

encode

encode(text: str) -> list[int]

text's ids, each added token it names its own id.

Source code in iron/lm/model.py

def encode(self, text: str) -> list[int]:
    """``text``'s ids, each added token it names its own id."""
    return self._tokenizer.encode(text, add_special_tokens=False).ids

prompt

prompt(text: str) -> list[int]

encode(text) after the special tokens the post-processor puts before a text.

A BOS, or none (Qwen's), as transformers' tokenizer(text) begins. Any it puts after a text are left off.

Source code in iron/lm/model.py

def prompt(self, text: str) -> list[int]:
    """``encode(text)`` after the special tokens the post-processor puts before a text.

    A BOS, or none (Qwen's), as transformers' ``tokenizer(text)`` begins.
    Any it puts after a text are left off.
    """
    encoded = self._tokenizer.encode(text)
    words = encoded.sequence_ids
    lead = words.index(0) if 0 in words else len(words)
    return encoded.ids[:lead] + self.encode(text)

token

token(id: int) -> str | None

Return the text of token id.

Source code in iron/lm/model.py

def token(self, id: int) -> str | None:
    """Return the text of token ``id``."""
    return self._tokenizer.id_to_token(id)

Stats dataclass

Stats(ttft: float, tokens_per_s: float, tokens: int)

A generation's timings.

The seconds to the first token, the tokens per second drawn after it (NaN for one token), and tokens, how many were drawn: those past the first eos too where they were drawn anyway.

Model

Model(config: Any, weights: Any)

Bases: Graph

A model compiled as several versions of one graph, one per input shape.

load builds it from a checkpoint's files, or from random weights at a seed (random_weights) for checks without real weights.

Source code in iron/lm/model.py

def __init__(self, config: Any, weights: Any) -> None: ...

hugging_face property

hugging_face: type

What loads the checkpoint into transformers to record its goldens.

A subclass needing another import (torch, transformers) only to record overrides this as a property of its own.

draws_at_once property

draws_at_once: bool

Whether, after a non-EOS first token, generate draws the rest before the second arrives.

Even those past eos cost no more to read.

eos cached property

eos: set[int]

The tokens generation stops at.

The checkpoint's eos_token_id, one or a list, from its generation_config.json or else its config.json; none without a checkpoint or either.

load classmethod

load(path, *, random_weights: int | None = None, max_seq_len: int | None = None, full_elf: bool | None = None) -> Self

from_checkpoint's model compiled and loaded, weights uploaded.

full_elf is every version's packaging (load_versions).

Source code in iron/lm/model.py

@classmethod
def load(
    cls,
    path,
    *,
    random_weights: int | None = None,
    max_seq_len: int | None = None,
    full_elf: bool | None = None,
) -> Self:
    """``from_checkpoint``'s model compiled and loaded, weights uploaded.

    ``full_elf`` is every version's packaging (``load_versions``).
    """
    model = cls.from_checkpoint(
        path, random_weights=random_weights, max_seq_len=max_seq_len
    )
    assert model.checkpoint is not None, "from_checkpoint gives it one"
    return model.load_versions(release=model.checkpoint.release, full_elf=full_elf)

from_checkpoint classmethod

from_checkpoint(path: str | Path, *, random_weights: int | None = None, max_seq_len: int | None = None) -> Self

Return the model of a Hugging Face repository's files, not yet loaded.

Its context is max_seq_len or else fitted to this host (fit).

Parameters:

Name	Type	Description	Default
path	str | Path	A directory holding config.json, weights and tokenizer.json, or a repository id to download them from.	required
random_weights	int | None	A seed: weights (and no tokenizer) drawn at it in the checkpoint's shapes instead of read.	None
max_seq_len	int | None	The longest context served, at most the checkpoint's own.	None

Raises:

Type	Description
ValueError	max_seq_len is past the checkpoint's own.

Source code in iron/lm/model.py

@classmethod
def from_checkpoint(
    cls,
    path: str | Path,
    *,
    random_weights: int | None = None,
    max_seq_len: int | None = None,
) -> Self:
    """Return the model of a Hugging Face repository's files, not yet loaded.

    Its context is ``max_seq_len`` or else fitted to this host (``fit``).

    Args:
        path: A directory holding ``config.json``, weights and
            ``tokenizer.json``, or a repository id to download them from.
        random_weights: A seed: weights (and no tokenizer) drawn at it in
            the checkpoint's shapes instead of read.
        max_seq_len: The longest context served, at most the
            checkpoint's own.

    Raises:
        ValueError: ``max_seq_len`` is past the checkpoint's own.
    """
    directory = Path(path)
    if not directory.is_dir():
        snapshot_download = _models_extra("huggingface_hub").snapshot_download
        files = [
            "*config.json",
            "model*.safetensors*",
            "tokenizer.json",
            "chat_template.jinja",
        ]
        directory = Path(snapshot_download(str(path), allow_patterns=files))
    weights = weights_file(directory)
    config = cls.default_config(weights)
    if max_seq_len is None:
        config = cls.fit(config, weights)
    else:
        assert isinstance(config, DecoderConfig), "a context is a decoder's"
        if max_seq_len > config.max_seq_len:
            raise ValueError(
                f"max_seq_len {max_seq_len} is past the model's "
                f"{config.max_seq_len}"
            )
        config = dataclasses.replace(config, max_seq_len=max_seq_len)
    checkpoint = Checkpoint(weights)
    tree = (
        checkpoint.tree()
        if random_weights is None
        else checkpoint.tree("random", random_weights)
    )
    model = cls(config, tree)
    model.checkpoint, model.random_weights = checkpoint, random_weights
    if random_weights is None:
        settings = directory / "tokenizer_config.json"
        jinja = directory / "chat_template.jinja"
        found = json.loads(settings.read_text()) if settings.exists() else {}
        # As transformers: chat_template.jinja over tokenizer_config.json's.
        template = (
            jinja.read_text() if jinja.exists() else found.get("chat_template")
        )
        # transformers renders tokenizer_config.json's bos_token, which
        # may be null or absent (nothing) or an added token's record.
        b = found.get("bos_token")
        bos_token = (b["content"] if isinstance(b, dict) else b) or ""
        encoding = TokenizerJson(directory / "tokenizer.json")
        model.tokenizer = Tokenizer(encoding, template, bos_token)
    return model

default_config classmethod

default_config(weights_path) -> Config

cls.config itself, if it is already one shape.

Else that class's own, built from weights_path's checkpoint directory (Config.from_checkpoint).

Source code in iron/lm/model.py

@classmethod
def default_config(cls, weights_path) -> Config:
    """``cls.config`` itself, if it is already one shape.

    Else that class's own, built from ``weights_path``'s checkpoint
    directory (``Config.from_checkpoint``).
    """
    if isinstance(cls.config, type):
        return cls.config.from_checkpoint(weights_path)
    return cls.config

fit classmethod

fit(config: Config, weights_path, budget: int | None = None) -> Config

config at its native context where the model's arena fits budget.

host_budget() by default, else at the longest whole number of prompt chunks that does, logged either way.

Source code in iron/lm/model.py

@classmethod
def fit(cls, config: Config, weights_path, budget: int | None = None) -> Config:
    """``config`` at its native context where the model's arena fits ``budget``.

    ``host_budget()`` by default, else at the longest whole number of
    prompt chunks that does, logged either way.
    """
    if not isinstance(config, DecoderConfig):
        return config
    native, C = config.max_seq_len, config.prefill_chunk
    if native < 2 * C:
        return config
    tree = Checkpoint(weights_path).tree("unread")
    held = []
    for length in (C, 2 * C):
        model = cls(dataclasses.replace(config, max_seq_len=length), tree)
        # What it holds does not follow its tunables, nor need a device.
        model.profile = None
        held.append(sum(model.record(**model.shapes()[0]).pinned.values()))
    per_position = (held[1] - held[0]) / C
    fixed = held[0] - per_position * C
    budget = host_budget() if budget is None else budget
    need = fixed + per_position * native
    if need <= budget:
        logger.info(
            "max_seq_len %d, the model's native context: %.1f GiB of a %.1f GiB "
            "budget",
            native,
            need / 2**30,
            budget / 2**30,
        )
        return config
    fits = int((budget - fixed) // per_position) // C * C
    if fits < C:
        raise ValueError(
            f"{fixed / 2**30:.1f} GiB of weights and one {C}-position chunk "
            f"do not fit a {budget / 2**30:.1f} GiB budget"
        )
    logger.warning(
        "max_seq_len %d, not the native %d: the model would hold %.1f GiB "
        "there, past this host's %.1f GiB budget (physical memory less %.0f "
        "GiB); --max-seq-len asks for up to %d",
        fits,
        native,
        need / 2**30,
        budget / 2**30,
        HOST_RESERVE / 2**30,
        native,
    )
    return dataclasses.replace(config, max_seq_len=fits)

cpu_reference

cpu_reference()

Return what the model is judged by.

With drawn weights, its own graph on the host (OnHost); else hugging_face's logits recorded beside the checkpoint (golden.npz), keyed by config.json and Checkpoint.fingerprint.

Source code in iron/lm/model.py

def cpu_reference(self):
    """Return what the model is judged by.

    With drawn weights, its own graph on the host (``OnHost``); else
    ``hugging_face``'s logits recorded beside the checkpoint
    (``golden.npz``), keyed by ``config.json`` and
    ``Checkpoint.fingerprint``.
    """
    if self.random_weights is not None:

        class Host(OnHost, type(self)):
            pass

        return Host(self.config, self.weights)
    assert self.checkpoint is not None, "a model without drawn weights loaded one"
    directory = Path(self.checkpoint.path).parent
    return Golden(
        directory / "golden.npz",
        lambda: self.hugging_face.load(directory),
        ((directory / "config.json").read_text(), self.checkpoint.fingerprint()),
    )

prompt

prompt(chars: int, skip: int = 0) -> list[int]

chars characters of prompt.txt from skip on, tokenized.

With no tokenizer, chars // 3 tokens drawn at the weights' seed and skip (roughly what the text would encode to).

Source code in iron/lm/model.py

def prompt(self, chars: int, skip: int = 0) -> list[int]:
    """``chars`` characters of ``prompt.txt`` from ``skip`` on, tokenized.

    With no tokenizer, ``chars // 3`` tokens drawn at the weights' seed
    and ``skip`` (roughly what the text would encode to).
    """
    if self.tokenizer is None:
        assert (
            self.random_weights is not None
        ), "without a tokenizer, weights drawn at random"
        rng = np.random.default_rng((self.random_weights, skip))
        return list(rng.integers(0, self.config.vocab_size, chars // 3))
    return self.tokenizer(prompt_text(chars, skip))

generate

generate(prompt: str | list[dict] | list[int], max_tokens: int, sampler: Sampler | None = None, **inputs: Any)

stream's pieces joined.

The text drawn after a text prompt or chat messages, else the token ids drawn.

Source code in iron/lm/model.py

def generate(
    self,
    prompt: str | list[dict] | list[int],
    max_tokens: int,
    sampler: Sampler | None = None,
    **inputs: Any,
):
    """``stream``'s pieces joined.

    The text drawn after a text prompt or chat messages, else the token
    ids drawn.
    """
    pieces = list(self._stream(prompt, max_tokens, sampler, inputs, False))
    return "".join(pieces) if _is_text(prompt) else pieces

stream

stream(prompt: str | list[dict] | list[int], max_tokens: int, sampler: Sampler | None = None, **inputs: Any) -> Iterator

Draw up to max_tokens after prompt.

Each yielded as it is drawn, through the first of eos, and its timings into stats once the last is drawn.

Parameters:

Name	Type	Description	Default
prompt	str | list[dict] | list[int]	Text or a list of chat messages (rendered by the checkpoint's chat template, Tokenizer.chat), yielded as text piece by piece; or token ids, yielded as ids.	required
max_tokens	int	The most tokens to draw.	required
sampler	Sampler | None	How each is drawn; None is greedy.	None
inputs	Any	What the model takes beside the tokens (Multimodal's raw).	{}

Raises:

Type	Description
RuntimeError	Text or messages, and no tokenizer.

Source code in iron/lm/model.py

def stream(
    self,
    prompt: str | list[dict] | list[int],
    max_tokens: int,
    sampler: Sampler | None = None,
    **inputs: Any,
) -> Iterator:
    """Draw up to ``max_tokens`` after ``prompt``.

    Each yielded as it is drawn, through the first of ``eos``, and its
    timings into ``stats`` once the last is drawn.

    Args:
        prompt: Text or a list of chat messages (rendered by the
            checkpoint's chat template, ``Tokenizer.chat``), yielded as
            text piece by piece; or token ids, yielded as ids.
        max_tokens: The most tokens to draw.
        sampler: How each is drawn; None is greedy.
        inputs: What the model takes beside the tokens
            (``Multimodal``'s ``raw``).

    Raises:
        RuntimeError: Text or messages, and no tokenizer.
    """
    return self._stream(prompt, max_tokens, sampler, inputs, True)

shapes

shapes() -> list[dict]

Return the input shapes of each version, the decode step (no input) first.

A later version on token ids (ids or placed) feeds off it where it is a full ELF, its Emit starting there.

Source code in iron/lm/model.py

def shapes(self) -> list[dict]:
    """Return the input shapes of each version, the decode step (no input) first.

    A later version on token ids (``ids`` or ``placed``) feeds off it
    where it is a full ELF, its Emit starting there.
    """
    raise NotImplementedError(f"{type(self).__name__} defines no shapes()")

load_versions

load_versions(*, release: Callable[[ndarray], None] | None = None, full_elf: bool | None = None, dev: Device | None = None) -> Self

Compile every version, then load each, in one arena, and hand them to loaded.

Every version is placed before the arena's buffers are made, so none grows them, which would copy them.

Parameters:

Name	Type	Description	Default
release	Callable[[ndarray], None] | None	Given each piece of each weight once it is on the device.	None
full_elf	bool | None	Each version's packaging (Graph.compile).	None
dev	Device | None	The device every version compiles for, else the current one.	None

Source code in iron/lm/model.py

def load_versions(
    self,
    *,
    release: Callable[[np.ndarray], None] | None = None,
    full_elf: bool | None = None,
    dev: Device | None = None,
) -> Self:
    """Compile every version, then load each, in one arena, and hand them to ``loaded``.

    Every version is placed before the arena's buffers are made, so none
    grows them, which would copy them.

    Args:
        release: Given each piece of each weight once it is on the device.
        full_elf: Each version's packaging (``Graph.compile``).
        dev: The device every version compiles for, else the current one.
    """
    if dev is None:
        dev = require_device()
    compiled, decode = [], None
    for shapes in self.shapes():
        prompt = "ids" in shapes or "placed" in shapes
        version = self.compile(
            dev=dev, feeds=decode if prompt else None, full_elf=full_elf, **shapes
        )
        if not shapes and version.emit is not None:
            decode = version
        compiled.append(version)
    try:
        for version in compiled:
            version.load(release=release)
        self.loaded(compiled)
    except BaseException:
        self.release()
        raise
    return self

loaded

loaded(versions: list) -> None

Run on versions, in shapes()'s order, once load_versions has them.

Nothing by default.

Source code in iron/lm/model.py

def loaded(self, versions: list) -> None:
    """Run on ``versions``, in ``shapes()``'s order, once ``load_versions`` has them.

    Nothing by default.
    """

OnHost

A model with its versions stood in by its graph's reference (Graph.reference).

The model with no NPU, mixed in before the model's own class; it compiles nothing, and generate draws on the host.

host_budget

host_budget() -> int

Return the bytes a model's arena may use on this host.

Physical memory minus HOST_RESERVE, or the cgroup's limit minus CGROUP_RESERVE when a cgroup caps this process, whichever is smaller.

Source code in iron/lm/model.py

def host_budget() -> int:
    """Return the bytes a model's arena may use on this host.

    Physical memory minus ``HOST_RESERVE``, or the cgroup's limit minus
    ``CGROUP_RESERVE`` when a cgroup caps this process, whichever is smaller.
    """
    phys = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") - HOST_RESERVE
    limit = _cgroup_limit()
    return phys if limit is None else min(phys, limit - CGROUP_RESERVE)

prompt_text

prompt_text(chars: int, skip: int = 0) -> str

chars characters of prompt.txt, from skip on, repeated past its end.

The text every model's long inputs are drawn from.

Source code in iron/lm/model.py

def prompt_text(chars: int, skip: int = 0) -> str:
    """``chars`` characters of ``prompt.txt``, from ``skip`` on, repeated past its end.

    The text every model's long inputs are drawn from.
    """
    text = Path(__file__).with_name("prompt.txt").read_text()
    return (text * -(-(skip + chars) // len(text)))[skip : skip + chars]

session

One conversation over a CausalLM's loaded versions.

The tokens its caches hold, and the calls that run a history's new tokens and draw after it. Several sessions share one compiled model; its caches hold one session's tokens at a time, so a session whose tokens another's calls replaced runs its history again from the first token.

Session

Session(lm)

A conversation over the versions of lm, a CausalLM (CausalLM.use).

Source code in iron/lm/session.py

def __init__(self, lm):
    self.lm = lm
    self._seen = NOTHING  # the tokens in the caches

logits

logits(tokens) -> ndarray

Return the logits after the last of tokens, the whole history.

One more token than the last call's is a decode step, anything else runs from the first token the caches do not hold (resume).

Raises:

Type	Description
ValueError	tokens is empty or longer than max_seq_len.
IndexError	A token is outside [0, vocab_size); refused before any dispatch, so the caches are as they were.

Source code in iron/lm/session.py

def logits(self, tokens) -> np.ndarray:
    """Return the logits after the last of ``tokens``, the whole history.

    One more token than the last call's is a decode step, anything else
    runs from the first token the caches do not hold (``resume``).

    Raises:
        ValueError: ``tokens`` is empty or longer than ``max_seq_len``.
        IndexError: A token is outside ``[0, vocab_size)``; refused
            before any dispatch, so the caches are as they were.
    """
    tokens = self._tokens(tokens)
    n, L = tokens.size, self.lm.config.max_seq_len
    if not 0 < n <= L:
        raise ValueError(f"{n} tokens do not fit {L} rows")
    self._take()
    held = self._held(tokens)
    if held == n - 1 == self._seen.size:
        out, _ = self.lm(token=int(tokens[-1]), position=n - 1)
    else:
        start = self.resume(held)
        self._seen = tokens[:start]
        for inputs, values in self._chunks(tokens, start):
            out, _ = self.lm(*inputs, **values)
    self._seen = tokens
    # A copy: the next call rewrites the output buffer. An odd vocabulary
    # is padded by a logit no token has.
    return np.array(out.numpy()).reshape(-1)[: self.lm.config.vocab_size]

draw

draw(tokens, max_tokens: int, sampler: Sampler, each: bool = True) -> Iterator[int]

Draw max_tokens after tokens on the device, each yielded as it is drawn.

As Model.stream does: looped there where the decode step is a full ELF (device_loop), else each step started from the host.

Raises:

Type	Description
ValueError	The tokens and those to draw do not fit the caches.
IndexError	A token is outside [0, vocab_size).
RuntimeError	Neither version is loaded: load_versions() first.

Source code in iron/lm/session.py

def draw(
    self, tokens, max_tokens: int, sampler: Sampler, each: bool = True
) -> Iterator[int]:
    """Draw ``max_tokens`` after ``tokens`` on the device, each yielded as it is drawn.

    As ``Model.stream`` does: looped there where the decode step is a
    full ELF (``device_loop``), else each step started from the host.

    Raises:
        ValueError: The tokens and those to draw do not fit the caches.
        IndexError: A token is outside ``[0, vocab_size)``.
        RuntimeError: Neither version is loaded: ``load_versions()``
            first.
    """
    model, tokens = self.lm, self._tokens(tokens)
    n, L = tokens.size, model.config.max_seq_len
    if not (0 < n and 0 < max_tokens and n + max_tokens - 1 <= L):
        raise ValueError(
            f"{n} tokens and {max_tokens} more to draw do not fit {L} rows"
        )
    if model.prompt_chunk is None or model.decode_step is None:
        raise RuntimeError(f"{type(model).__name__}: load_versions() first")
    self._take()
    # The prompt draws at its last position, n - 1; decode step k at n + k.
    rows = np.zeros((L, ROW_WORDS), dtype=np.int32)
    k_max = next(
        s.op.k_max
        for s in model.decode_step.recording.steps
        if isinstance(s.op, Sample)
    )
    rows[n - 1 : n - 1 + max_tokens] = sampler.rows(max_tokens, k_max)
    model.decode_step.write(model.draws, rows)
    start = self.resume(self._held(tokens))
    # A caller may stop early, leaving steps in the caches nothing records:
    # until the last token, they hold nothing.
    self._seen = NOTHING
    drawn = []
    chunks = self._chunks(tokens, start)
    if model.device_loop:
        draws = device_draws(model, chunks, n, max_tokens, each)
    else:
        draws = host_draws(model, chunks, max_tokens)
    for token in draws:
        drawn.append(token)
        yield token
    # The caches hold every token but the last drawn, which no step read.
    self._seen = np.concatenate([tokens, drawn[:-1]]).astype(np.int64)

forget

forget(position: int) -> None

Hold nothing the caches have from position on.

The next call runs from there or an earlier restart boundary, as a change before it (a placeholder's soft tokens) would leave them stale.

Source code in iron/lm/session.py

def forget(self, position: int) -> None:
    """Hold nothing the caches have from ``position`` on.

    The next call runs from there or an earlier restart boundary, as a
    change before it (a placeholder's soft tokens) would leave them
    stale.
    """
    held = self._seen[:position].size
    self._seen = self._seen[: self.resume(held)]

resume

resume(held: int) -> int

Return the position calls start from to run the tokens from held on.

held itself, or for a state that cannot restart there (KVCache.restart), the earliest one of them can.

Source code in iron/lm/session.py

def resume(self, held: int) -> int:
    """Return the position calls start from to run the tokens from ``held`` on.

    ``held`` itself, or for a state that cannot restart there
    (``KVCache.restart``), the earliest one of them can.
    """
    ran = self._seen.size
    return min([held, *(c.restart(held, ran) for c in self.lm.recurrent)])

prompt_ids

prompt_ids(tokens) -> ndarray

Return a prompt call's ids for a chunk of tokens, padded to prefill_chunk.

Padded with row 0, which the bound by rows leaves unread.

Source code in iron/lm/session.py

def prompt_ids(self, tokens) -> np.ndarray:
    """Return a prompt call's ``ids`` for a chunk of ``tokens``, padded to ``prefill_chunk``.

    Padded with row 0, which the bound by ``rows`` leaves unread.
    """
    tokens = np.asarray(tokens)
    ids = np.zeros(self.lm.config.prefill_chunk, dtype=np.int32)
    ids[: tokens.size] = tokens
    return ids

host_draws

host_draws(model, chunks, max_tokens: int) -> Iterator[int]

Yield max_tokens drawn after the prompt chunks.

The host starts each decode step on the token the last one drew, which it reads alone.

Source code in iron/lm/session.py

def host_draws(model, chunks, max_tokens: int) -> Iterator[int]:
    """Yield ``max_tokens`` drawn after the prompt ``chunks``.

    The host starts each decode step on the token the last one drew, which
    it reads alone.
    """
    for inputs, values in chunks:
        _, carry = model(*inputs, **values)
    yield int(carry["token"])
    for _ in range(max_tokens - 1):
        _, carry = model(**carry)
        yield int(carry["token"])

device_draws

device_draws(model, chunks, n: int, max_tokens: int, each: bool) -> Iterator[int]

Yield max_tokens drawn after the prompt chunks of n tokens.

The last chunk starts a device loop of decode steps (CausalLM.loop); with each, each token is read once its step completes.

Source code in iron/lm/session.py

def device_draws(model, chunks, n: int, max_tokens: int, each: bool) -> Iterator[int]:
    """Yield ``max_tokens`` drawn after the prompt ``chunks`` of ``n`` tokens.

    The last chunk starts a device loop of decode steps (``CausalLM.loop``);
    with ``each``, each token is read once its step completes.
    """
    for inputs, values in chunks:
        if values["position"] < n - 1:
            model(*inputs, **values)
            continue
        loop = model.loop(
            model.prompt_chunk if len(inputs) == 1 else model.placed_chunk
        )
        loop.start(*inputs, **values)
    read = partial(model.decode_step.read, model.drawn)
    yield int(read()[n - 1])
    for k, _ in enumerate(loop.steps(max_tokens - 1), start=1):
        if each:
            yield int(read()[n - 1 + k])
    if not each:
        yield from (int(t) for t in read()[n : n - 1 + max_tokens])

generation

Sampling, and the loops that generate and check, over any model.

A model is anything with logits(tokens): the logits after the last of tokens, the whole history so far, as (vocab_size,). A CausalLM on the NPU is one, and so is a CPU reference.

Sampler

Sampler(temperature: float = 0.7, top_k: int | None = 50, seed: int = SEED)

Temperature, then top-k, then a draw from the softmax, exactly as the device draws.

aie.iron.kernels.sample.sample_ref, so the host and device pick the same token from the same logits and uniform.

Parameters:

Name	Type	Description	Default
top_k	int | None	None keeps the whole row; only valid on the host, since the device's top-k is bounded.	50

Source code in iron/lm/generation.py

def __init__(
    self,
    temperature: float = 0.7,
    top_k: int | None = 50,
    seed: int = SEED,
):
    if temperature < 0:
        raise ValueError(f"temperature {temperature} is negative")
    if temperature > 0:
        # The draw thresholds bf16 keys, which is exact only when
        # dividing by the temperature keeps every logit's order.
        check_order_preserving(temperature)
    if top_k is not None and top_k < 1:
        raise ValueError(f"top_k {top_k} keeps no token")
    self.temperature = np.float32(temperature)
    self.top_k = top_k
    self.rng = np.random.default_rng(seed)

probabilities

probabilities(logits: ndarray) -> ndarray

Return the distribution a token is drawn from, float64, over logits (1-D).

Source code in iron/lm/generation.py

def probabilities(self, logits: np.ndarray) -> np.ndarray:
    """Return the distribution a token is drawn from, float64, over ``logits`` (1-D)."""
    size = np.asarray(logits).size
    probs = np.zeros(size, dtype=np.float64)
    if self.temperature == 0:
        probs[sample_ref(logits, self.temperature, 1, 0)] = 1.0
        return probs
    candidates, weights = sample_weights(
        logits, self.temperature, self._top_k(size)
    )
    probs[candidates] = weights / weights.sum()
    return probs

rows

rows(count: int, k_max: int) -> ndarray

count draws for the device, (count, ROW_WORDS) int32, one per position.

Taking the uniforms count host draws would, in order.

Raises:

Type	Description
ValueError	top_k is None or past k_max, the device's largest, where it would be clamped rather than honoured.

Source code in iron/lm/generation.py

def rows(self, count: int, k_max: int) -> np.ndarray:
    """``count`` draws for the device, ``(count, ROW_WORDS)`` int32, one per position.

    Taking the uniforms ``count`` host draws would, in order.

    Raises:
        ValueError: ``top_k`` is None or past ``k_max``, the device's
            largest, where it would be clamped rather than honoured.
    """
    if self.top_k is None or self.top_k > k_max:
        raise ValueError(
            f"top_k {self.top_k} is not one the device draws (1..{k_max})"
        )
    rows = np.empty((count, ROW_WORDS), dtype=np.int32)
    for i in range(count):
        rows[i] = draw_row(self.temperature, self.top_k, self._n53())
    return rows

greedy

greedy(logits: ndarray) -> int

Return the id of the largest of logits.

Source code in iron/lm/generation.py

def greedy(logits: np.ndarray) -> int:
    """Return the id of the largest of ``logits``."""
    return int(np.argmax(np.asarray(logits, dtype=np.float32).reshape(-1)))

generate_on_host

generate_on_host(model, tokens, num_tokens, sample) -> list[int]

Return the num_tokens drawn after tokens, each from model's logits.

Source code in iron/lm/generation.py

def generate_on_host(model, tokens, num_tokens, sample) -> list[int]:
    """Return the ``num_tokens`` drawn after ``tokens``, each from ``model``'s logits."""
    history = [int(t) for t in tokens]
    for _ in range(num_tokens):
        history.append(sample(model.logits(history)))
    return history[-num_tokens:]

divergence

divergence(reference, logits) -> tuple[float, bool]

KL(reference || model) of two rows of logits, and whether both rank the same token first.

Source code in iron/lm/generation.py

def divergence(reference, logits) -> tuple[float, bool]:
    """KL(reference || model) of two rows of logits, and whether both rank the same token first."""
    ref, got = _log_softmax(reference), _log_softmax(logits)
    kl = float(np.sum(np.exp(ref) * (ref - got)))
    return kl, int(got.argmax()) == int(ref.argmax())

accuracy

accuracy(model, reference, tokens, num_tokens) -> list[tuple[float, bool]]

model's next-token distributions against reference's, over num_tokens steps from tokens.

Teacher-forced on the reference's greedy token so a divergence is the model's own error, not an earlier different choice's consequence.

Source code in iron/lm/generation.py

def accuracy(model, reference, tokens, num_tokens) -> list[tuple[float, bool]]:
    """``model``'s next-token distributions against ``reference``'s, over ``num_tokens`` steps from ``tokens``.

    Teacher-forced on the reference's greedy token so a divergence is the
    model's own error, not an earlier different choice's consequence.
    """
    history, results = [int(t) for t in tokens], []
    for step in range(num_tokens):
        ref = reference.logits(history)
        kl, top1 = divergence(ref, model.logits(history))
        results.append((kl, top1))
        print(f"step {step:3d}  KL {kl:.5f}  top-1 {'match' if top1 else 'MISMATCH'}")
        history.append(greedy(ref))
    return results

kl_stats

kl_stats(results) -> dict[str, float]

Return the mean, p90 and max KL of accuracy's results, over every step, prefill and decode alike.

One step's KL depends as much on how confident the reference is there as on the model.

Source code in iron/lm/generation.py

def kl_stats(results) -> dict[str, float]:
    """Return the mean, p90 and max KL of ``accuracy``'s results, over every step, prefill and decode alike.

    One step's KL depends as much on how confident the reference is there
    as on the model.
    """
    kl = np.array([k for k, _ in results])
    return {"Mean": kl.mean(), "P90": np.percentile(kl, 90), "Max": kl.max()}

greedy_logits

greedy_logits(model, tokens, num_tokens) -> ndarray

Return the logits of num_tokens greedy steps from tokens, one row each.

The first the prompt's.

Source code in iron/lm/generation.py

def greedy_logits(model, tokens, num_tokens) -> np.ndarray:
    """Return the logits of ``num_tokens`` greedy steps from ``tokens``, one row each.

    The first the prompt's.
    """
    history, rows = [int(t) for t in tokens], []
    for _ in range(num_tokens):
        rows.append(model.logits(history))
        history.append(greedy(rows[-1]))
    return np.stack(rows)

differing_steps

differing_steps(run, first) -> list[int]

Return the steps at which two greedy_logits runs differ bitwise, NaNs and signed zeros included.

Source code in iron/lm/generation.py

def differing_steps(run, first) -> list[int]:
    """Return the steps at which two ``greedy_logits`` runs differ bitwise, NaNs and signed zeros included."""
    run, first = (np.ascontiguousarray(r).view(np.uint8) for r in (run, first))
    return np.flatnonzero((run != first).any(axis=1)).tolist()

determinism

determinism(model, prompts, num_tokens, rounds) -> int

How many runs' logits differ bitwise from the first of their prompt, over rounds repeats of each prompt.

Alternating so a host write that never reaches the device reads the other prompt's data, not a leftover copy of its own.

Source code in iron/lm/generation.py

def determinism(model, prompts, num_tokens, rounds) -> int:
    """How many runs' logits differ bitwise from the first of their prompt, over ``rounds`` repeats of each prompt.

    Alternating so a host write that never reaches the device reads the
    other prompt's data, not a leftover copy of its own.
    """
    first: list = [None] * len(prompts)
    differ = 0
    for r in range(rounds * len(prompts)):
        p = r % len(prompts)
        run = greedy_logits(model, prompts[p], num_tokens)
        if first[p] is None:
            first[p] = run
            continue
        steps = differing_steps(run, first[p])
        if steps:
            differ += 1
            print(f"round {r} (prompt {p}): logits differ at steps {steps}")
    return differ

golden

What a model is judged by.

Hugging Face transformers' own model in float32 (HuggingFace), its logits recorded once to a file (Golden) so a test or an accuracy run needs neither torch nor transformers, which are not IRON's requirements, only to record.

Golden

Golden(path: str | Path, reference: Callable[[], Any] | None = None, inputs: tuple[Any, ...] = (), *, record: bool | None = None)

Logits by token sequence and any arrays beside it (a clip, an image).

Each is the logits after a sequence's last token, read from path (an .npz).

Parameters:

Name	Type	Description	Default
path	str | Path	The file; it need not exist yet.	required
reference	Callable[[], Any] | None	Builds what records a sequence path lacks (a HuggingFace), called once on the first such sequence, and only while recording; otherwise such a sequence raises LookupError, or FileNotFoundError naming path where there is no file.	None
inputs	tuple[Any, ...]	What the reference is built from (a config, weights); the file holds their fingerprint and is refused if they change.	()
record	bool | None	Whether to record what path lacks; None while RECORD is set.	None

Raises:

Type	Description
ValueError	path was recorded from other inputs.

Source code in iron/lm/golden.py

def __init__(
    self,
    path: str | Path,
    reference: Callable[[], Any] | None = None,
    inputs: tuple[Any, ...] = (),
    *,
    record: bool | None = None,
):
    self.path = Path(path)
    if record is None:
        record = bool(os.environ.get(RECORD))
    self.build = reference if record else None
    self.reference = None
    self.fingerprint = fingerprint(*inputs) if inputs else None
    self.rows = self.read()

read

read() -> dict

Return the file's rows; none if it is stale and this records.

Source code in iron/lm/golden.py

def read(self) -> dict:
    """Return the file's rows; none if it is stale and this records."""
    if not self.path.exists():
        return {}
    with np.load(self.path) as f:
        rows = dict(f)
    held = str(rows.pop("fingerprint", ""))
    if self.fingerprint is None or held == self.fingerprint:
        return rows
    if self.build is not None:
        return {}
    raise ValueError(
        f"{self.path} was recorded from other inputs: re-record it where "
        "transformers is installed (pytest --record-goldens)"
    )

logits

logits(tokens, *arrays) -> ndarray

Return the logits after tokens given arrays.

The reference's logits takes arrays after tokens.

Source code in iron/lm/golden.py

def logits(self, tokens, *arrays) -> np.ndarray:
    """Return the logits after ``tokens`` given ``arrays``.

    The reference's ``logits`` takes ``arrays`` after ``tokens``.
    """
    tokens = np.asarray(tokens, dtype=np.int64).reshape(-1)
    digest = hashlib.sha256(tokens.tobytes())
    for array in map(np.ascontiguousarray, arrays):
        digest.update(f"{array.dtype}{array.shape}".encode() + array.tobytes())
    key = digest.hexdigest()[:16]
    if key not in self.rows:
        if self.build is None and not self.path.exists():
            raise FileNotFoundError(
                f"{self.path} is missing: record it where torch and "
                "transformers are installed (pytest --record-goldens) and "
                "install it at that path"
            )
        if self.build is None:
            raise LookupError(
                f"{self.path} holds no logits for these {tokens.size} tokens: "
                "record them where transformers is installed (pytest "
                "--record-goldens)"
            )
        if self.reference is None:
            self.reference = self.build()
        # A copy: a reference's row may view its whole forward's logits.
        row = np.array(self.reference.logits(tokens, *arrays), np.float32)
        # Another process may have recorded others since this one read.
        self.rows = self.read() | self.rows | {key: row}
        self.write()
    return self.rows[key]

write

write()

Write the rows to the file, replaced whole so no reader sees half.

Source code in iron/lm/golden.py

def write(self):
    """Write the rows to the file, replaced whole so no reader sees half."""
    self.path.parent.mkdir(parents=True, exist_ok=True)
    held: dict[str, Any] = (
        {} if self.fingerprint is None else {"fingerprint": self.fingerprint}
    )
    partial = self.path.with_name(f".{self.path.stem}.{os.getpid()}.npz")
    np.savez_compressed(partial, **self.rows, **held)
    os.replace(partial, self.path)

HuggingFace

HuggingFace(model)

transformers' own model in float32, of the auto class model_class names.

A causal model's by default: a checkpoint directory's (load) or one built from an IRON config and weights (build).

Source code in iron/lm/golden.py

def __init__(self, model):
    self.model = model.eval()

build classmethod

build(model_type: str, config, weights, **json) -> HuggingFace

Build model_type's model on config's fields as its config.json names them.

json's merged in, over weights written to a checkpoint and loaded, so transformers renames (and stacks) them as it loads one.

Raises:

Type	Description
ValueError	weights lacks a parameter the model holds, other than one tied to another, or holds one it does not.

Source code in iron/lm/golden.py

@classmethod
def build(cls, model_type: str, config, weights, **json) -> "HuggingFace":
    """Build ``model_type``'s model on ``config``'s fields as its ``config.json`` names them.

    ``json``'s merged in, over ``weights`` written to a checkpoint
    and loaded, so transformers renames (and stacks) them as it
    loads one.

    Raises:
        ValueError: ``weights`` lacks a parameter the model holds, other
            than one tied to another, or holds one it does not.
    """
    import torch
    import transformers
    from safetensors.torch import save_file

    hf = transformers.AutoConfig.for_model(
        model_type, **_merged(config.to_json(), json)
    )
    auto = getattr(transformers, cls.model_class)
    with tempfile.TemporaryDirectory() as directory:
        hf.save_pretrained(directory)
        state = {k: torch.from_numpy(v) for k, v in flat(weights)}
        save_file(state, f"{directory}/model.safetensors")
        model, info = auto.from_pretrained(
            directory, dtype=torch.float32, output_loading_info=True
        )
    missing, unexpected = info["missing_keys"], info["unexpected_keys"]
    if missing or unexpected:
        raise ValueError(
            f"{model_type}: the weights lack {sorted(missing)} and hold "
            f"{sorted(unexpected)}, which the model does not"
        )
    return cls(model)

logits

logits(tokens, *arrays: Any, **kwargs: Any) -> ndarray

Return the logits after the last of tokens, float32.

A subclass whose model takes more than token ids (an image, a clip) reads them from arrays.

Source code in iron/lm/golden.py

def logits(self, tokens, *arrays: Any, **kwargs: Any) -> np.ndarray:
    """Return the logits after the last of ``tokens``, float32.

    A subclass whose model takes more than token ids (an image, a clip)
    reads them from ``arrays``.
    """
    import torch

    with torch.no_grad():
        out = self.model(torch.tensor([list(tokens)]), use_cache=False)
    return out.logits[0, -1].numpy()

fingerprint

fingerprint(*values) -> str

Return a digest of values, a weights tree's leaves by name (leaves).

Arrays by their bytes, shapes and types, anything else by repr.

Source code in iron/lm/golden.py

def fingerprint(*values) -> str:
    """Return a digest of ``values``, a weights tree's leaves by name (``leaves``).

    Arrays by their bytes, shapes and types, anything else by ``repr``.
    """
    from .checkpoint import leaves  # which imports this module

    digest = hashlib.sha256()
    for name, value in leaves(list(values)):
        digest.update(name.encode())
        if isinstance(value, np.ndarray):
            digest.update(f"{value.dtype}{value.shape}".encode())
            digest.update(np.ascontiguousarray(value).tobytes())
        else:
            digest.update(_text(value).encode())
    return digest.hexdigest()[:16]

flat

flat(tree)

tree's arrays by their dotted names, in float32.

Source code in iron/lm/golden.py

def flat(tree):
    """``tree``'s arrays by their dotted names, in float32."""
    from .checkpoint import leaves  # which imports this module

    return ((name, value.astype(np.float32)) for name, value in leaves(tree))

testing

What a model's device test checks: speed, accuracy against the CPU reference, and determinism.

Each over a model Model.load loaded, its figures through the test's record_property; and where it finds the model's files.

Download dataclass

Download(url: str, sha256: str)

A file a test reads that the repository does not hold, fetched from url and pinned by its sha256.

fetch

fetch() -> Path

Return the file in cache_dir(), by digest.

Fetched unless a file of the pinned content is already there.

Raises:

Type	Description
ValueError	url is not https.
RuntimeError	What url serves is not the pinned content.

Source code in iron/lm/testing.py

def fetch(self) -> Path:
    """Return the file in ``cache_dir()``, by digest.

    Fetched unless a file of the pinned content is already there.

    Raises:
        ValueError: ``url`` is not https.
        RuntimeError: What ``url`` serves is not the pinned content.
    """
    target = cache_dir() / self.url.rsplit("/", 1)[-1]
    if target.exists():
        with open(target, "rb") as f:
            if hashlib.file_digest(f, "sha256").hexdigest() == self.sha256:
                return target
    if not self.url.startswith("https://"):
        raise ValueError(f"refusing to download over {self.url!r}")
    target.parent.mkdir(parents=True, exist_ok=True)
    # Beside the target and renamed, so an interrupted fetch leaves no
    # truncated file behind; by process, so concurrent workers share none.
    partial = target.with_name(f".{target.name}.{os.getpid()}.part")
    with urllib.request.urlopen(self.url, timeout=60) as response:
        data = response.read()
    if (got := hashlib.sha256(data).hexdigest()) != self.sha256:
        raise RuntimeError(f"{self.url} has SHA-256 {got}, expected {self.sha256}")
    partial.write_bytes(data)
    partial.replace(target)
    return target

weights_dir

weights_dir(name: str) -> Path

Where the tests find a model's files: $IRON_EXAMPLE_WEIGHTS_DIR/<name>.

Source code in iron/lm/testing.py

def weights_dir(name: str) -> Path:
    """Where the tests find a model's files: ``$IRON_EXAMPLE_WEIGHTS_DIR/<name>``."""
    return Path(os.environ.get("IRON_EXAMPLE_WEIGHTS_DIR", "/srv")) / name

cache_dir

cache_dir() -> Path

Where a Download is cached.

IRON's own writable cache (NPU_CACHE_HOME), not the weights directory, which CI mounts read-only.

Source code in iron/lm/testing.py

def cache_dir() -> Path:
    """Where a ``Download`` is cached.

    IRON's own writable cache (``NPU_CACHE_HOME``), not the weights
    directory, which CI mounts read-only.
    """
    return Path(NPU_CACHE_HOME) / "data"

require

require(*files: Path, staged: bool = True) -> None

Skip the test unless every one of files exists, except in CI where they are staged.

There a missing file fails it, naming what the runner must hold.

Source code in iron/lm/testing.py

def require(*files: Path, staged: bool = True) -> None:
    """Skip the test unless every one of ``files`` exists, except in CI where they are ``staged``.

    There a missing file fails it, naming what the runner must hold.
    """
    missing = [str(f) for f in files if not f.exists()]
    if not missing:
        return
    import pytest  # a test's: a model imports this module for weights_dir

    if staged and os.environ.get("CI"):
        pytest.fail(
            f"this CI runner lacks {', '.join(missing)}: install the model's "
            "files at that path; its README says which and how",
            pytrace=False,
        )
    pytest.skip(f"not found outside CI: {missing}")

load_from_flags

load_from_flags(request, model, weights: Path, tokenizer='tokenizer.json', staged=True)

model.load of weights' directory, to call with any further options.

As the CLI loads it at pytest's --random-weights and --max-seq-len; requires the files it needs.

Source code in iron/lm/testing.py

def load_from_flags(
    request, model, weights: Path, tokenizer="tokenizer.json", staged=True
):
    """``model.load`` of ``weights``' directory, to call with any further options.

    As the CLI loads it at pytest's ``--random-weights`` and
    ``--max-seq-len``; ``require``s the files it needs.
    """
    seed = request.config.getoption("--random-weights")
    files = [weights, weights.with_name("config.json")]
    if seed is None:
        files.append(weights.with_name(tokenizer))
    require(*files, staged=staged)
    max_seq_len = request.config.getoption("--max-seq-len")
    return functools.partial(
        model.load, weights.parent, random_weights=seed, max_seq_len=max_seq_len
    )

prompt

prompt(model, chars: int, num_tokens: int, skip: int = 0) -> list[int]

Return a prompt of chars characters' tokens with room for num_tokens after it.

Source code in iron/lm/testing.py

def prompt(model, chars: int, num_tokens: int, skip: int = 0) -> list[int]:
    """Return a prompt of ``chars`` characters' tokens with room for ``num_tokens`` after it."""
    tokens = model.prompt(chars, skip=skip)
    assert len(tokens) + num_tokens <= model.config.max_seq_len
    return tokens

check_generation

check_generation(model, prompt_len: int, num_tokens: int, *, record)

Sample num_tokens on the device after a prompt of prompt_len characters.

CausalLM.generate, the whole prompt prefilled; record the time to the first token and the tokens per second after it.

Source code in iron/lm/testing.py

def check_generation(model, prompt_len: int, num_tokens: int, *, record):
    """Sample ``num_tokens`` on the device after a prompt of ``prompt_len`` characters.

    ``CausalLM.generate``, the whole prompt prefilled; record the time to
    the first token and the tokens per second after it.
    """
    tokens = prompt(model, prompt_len, num_tokens)
    # A module's model holds the last test's tokens: its time is the prompt's.
    model.forget(0)
    drawn = model.generate(tokens, num_tokens, Sampler())
    print(drawn if model.tokenizer is None else model.tokenizer.decode(drawn))
    record("TTFT", model.stats.ttft)
    if num_tokens > 1:
        record("TPS", model.stats.tokens_per_s)
    return tokens, drawn

check_prompt_time

check_prompt_time(model, length: int, runs: int = 3, *, record)

Record the median seconds to the first token after a prompt of length tokens.

Over runs timed runs, each a fresh prompt (random ids, since a prompt's time does not follow its ids).

Source code in iron/lm/testing.py

def check_prompt_time(model, length: int, runs: int = 3, *, record):
    """Record the median seconds to the first token after a prompt of ``length`` tokens.

    Over ``runs`` timed runs, each a fresh prompt (random ids, since a
    prompt's time does not follow its ids).
    """
    rng = np.random.default_rng(SEED)
    seconds = []
    for _ in range(runs):
        tokens = list(rng.integers(0, model.config.vocab_size, length))
        model.generate(tokens, 1, Sampler())
        seconds.append(model.stats.ttft)
    record("TTFT", float(np.median(seconds)))
    record("Tokens", length)

check_device_sampling

check_device_sampling(model, prompt_len: int, num_tokens: int, *, record)

Sample num_tokens on the device (CausalLM.generate) after a prompt of prompt_len characters.

Then sample again on the host from each step's logits at the same seed: the text must match, token for token.

Source code in iron/lm/testing.py

def check_device_sampling(model, prompt_len: int, num_tokens: int, *, record):
    """Sample ``num_tokens`` on the device (``CausalLM.generate``) after a prompt of ``prompt_len`` characters.

    Then sample again on the host from each step's logits at the same seed:
    the text must match, token for token.
    """
    tokens, drawn = check_generation(model, prompt_len, num_tokens, record=record)
    host = generate_on_host(model, tokens, num_tokens, Sampler())
    differing = [i for i, (a, b) in enumerate(zip(drawn, host)) if a != b]
    record("DifferingTokens", len(differing))
    assert not differing, f"tokens {differing} of {num_tokens} differ from the host's"

assert_close

assert_close(got, want: ndarray, what: str = 'the logits', *, share: float = LOGITS) -> None

got within share of want, the reference's logits, at their largest magnitude.

And, where the reference's top two differ by more than that bound, the same argmax.

Source code in iron/lm/testing.py

def assert_close(
    got, want: np.ndarray, what: str = "the logits", *, share: float = LOGITS
) -> None:
    """``got`` within ``share`` of ``want``, the reference's logits, at their largest magnitude.

    And, where the reference's top two differ by more than that bound, the
    same argmax.
    """
    got = np.asarray(got, np.float32).reshape(want.shape)
    scale = np.abs(want).max()
    bound, err = share * scale, np.abs(got - want).max()
    assert err <= bound, f"{what}: max |diff| {err:.4f} against |logits| {scale:.3f}"
    second, first = np.sort(want, axis=None)[-2:]
    if first - second > bound:
        assert (
            got.argmax() == want.argmax()
        ), f"{what}: argmax {got.argmax()} != {want.argmax()}"

check_accuracy

check_accuracy(model, bounds: dict[str, float], num_tokens: int = 40, chars: int = 1024, *, record)

KL(reference || model) of the next-token distribution, teacher-forced over num_tokens steps.

From a chars-character prompt, within bounds ({"Mean": .., "P90": .., "Max": ..}).

Source code in iron/lm/testing.py

def check_accuracy(
    model,
    bounds: dict[str, float],
    num_tokens: int = 40,
    chars: int = 1024,
    *,
    record,
):
    """KL(reference || model) of the next-token distribution, teacher-forced over ``num_tokens`` steps.

    From a ``chars``-character prompt, within ``bounds``
    (``{"Mean": .., "P90": .., "Max": ..}``).
    """
    tokens = prompt(model, chars, num_tokens)
    model.forget(0)
    results = accuracy(model, model.cpu_reference(), tokens, num_tokens)
    stats = kl_stats(results)
    for stat, value in stats.items():
        record(f"{stat}KL", float(value))
    record("Top1Mismatches", sum(not t for _, t in results))
    for stat, bound in bounds.items():
        assert stats[stat] <= bound, f"{stat.lower()} KL {stats[stat]} > {bound}"

check_embedding

check_embedding(got: ndarray, want: ndarray, min_cosine: float, max_error: float, *, percentile: float = 1, record: Callable[[str, float], object])

got's rows against want's embeddings.

The token cosines' low percentile (one bad token should not fail every other) and the relative error of the whole array.

Parameters:

Name	Type	Description	Default
got	ndarray	The array judged.	required
want	ndarray	The reference array got is judged against.	required
min_cosine	float	The percentile cosine must reach.	required
max_error	float	||got - want|| / ||want|| must not exceed.	required
percentile	float	Which percentile of the per-row cosines is judged (1: every row but the worst may miss min_cosine).	1
record	Callable[[str, float], object]	A test's record_property: the judged cosine, the worst row's and the relative error.	required

Source code in iron/lm/testing.py

def check_embedding(
    got: np.ndarray,
    want: np.ndarray,
    min_cosine: float,
    max_error: float,
    *,
    percentile: float = 1,
    record: Callable[[str, float], object],
):
    """``got``'s rows against ``want``'s embeddings.

    The token cosines' low percentile (one bad token should not fail every
    other) and the relative error of the whole array.

    Args:
        got: The array judged.
        want: The reference array ``got`` is judged against.
        min_cosine: The percentile cosine must reach.
        max_error: ``||got - want|| / ||want||`` must not exceed.
        percentile: Which percentile of the per-row cosines is judged (1:
            every row but the worst may miss ``min_cosine``).
        record: A test's ``record_property``: the judged cosine, the worst
            row's and the relative error.
    """
    cosines = (got * want).sum(-1) / (
        np.linalg.norm(got, axis=-1) * np.linalg.norm(want, axis=-1)
    )
    error = float(np.linalg.norm(got - want) / np.linalg.norm(want))
    judged = float(np.percentile(cosines, percentile))
    record("Cosine", judged)
    record("Worst cosine", float(cosines.min()))
    record("Relative error", error)
    assert judged >= min_cosine, judged
    assert error <= max_error, error

check_determinism

check_determinism(model, num_tokens: int = 4, rounds: int = 5, chars: int = 1024, *, record)

Repeated runs produce bit-identical logits, over two prompts of chars characters.

Alternated so a host write that never reaches the device reads the other prompt's data, failing every run rather than some.

Source code in iron/lm/testing.py

def check_determinism(
    model,
    num_tokens: int = 4,
    rounds: int = 5,
    chars: int = 1024,
    *,
    record,
):
    """Repeated runs produce bit-identical logits, over two prompts of ``chars`` characters.

    Alternated so a host write that never reaches the device reads the
    other prompt's data, failing every run rather than some.
    """
    prompts = [
        prompt(model, chars, num_tokens),
        prompt(model, chars, num_tokens, skip=chars),
    ]
    differing = determinism(model, prompts, num_tokens, rounds)
    record("DifferingRuns", differing)
    total = len(prompts) * (rounds - 1)
    assert differing == 0, f"{differing}/{total} runs' logits differ bitwise"

check_chat_turn

check_chat_turn(model, chars: int, turn: int, num_tokens: int = 2, *, record)

Extend a prompt by a turn, as a chat's next message extends it.

The model runs from the turn's first token (CausalLM.resume), and its logits match bit for bit a fresh run of the whole prompt from scratch.

Source code in iron/lm/testing.py

def check_chat_turn(model, chars: int, turn: int, num_tokens: int = 2, *, record):
    """Extend a prompt by a turn, as a chat's next message extends it.

    The model runs from the turn's first token (``CausalLM.resume``), and
    its logits match bit for bit a fresh run of the whole prompt from scratch.
    """
    whole = prompt(model, chars, num_tokens)
    C = model.config.prefill_chunk
    assert (len(whole) - turn) // C, "the turn reuses no chunk"
    other = prompt(model, chars, 0, skip=1)
    runs = []
    for name, before in (("Turn", whole[:-turn]), ("Scratch", other)):
        model.logits(before)
        if name == "Scratch":
            model.session.forget(0)  # other shares the turn's first token
        start = time.perf_counter()
        runs.append(greedy_logits(model, whole, num_tokens))
        record(f"{name}Seconds", time.perf_counter() - start)
    steps = differing_steps(*runs)
    assert not steps, f"the turn's logits differ from scratch's at steps {steps}"

check_deep_decode

check_deep_decode(model, position: int, bound: float, chars: int = 1024, *, record)

Check a decode step at position, deep in the caches, against the graph's own reference.

Graph.reference, for a context transformers would take too long over, within bound of its KL.

Source code in iron/lm/testing.py

def check_deep_decode(model, position: int, bound: float, chars: int = 1024, *, record):
    """Check a decode step at ``position``, deep in the caches, against the graph's own reference.

    ``Graph.reference``, for a context transformers would take too long
    over, within ``bound`` of its KL.
    """
    tokens = prompt(model, chars, 1)
    token = greedy(model.logits(tokens))
    # A decode step takes the token; the device gathers its embedding row.
    values = dict(token=token, position=position)
    with _fill_caches(model, len(tokens), position):
        got, _ = model(**values)
        got = got.numpy()
        expected, _ = model.reference(**values)
    kl, top1 = divergence(expected, got)
    record("DeepKL", kl)
    record("DeepTop1", top1)
    assert kl <= bound, f"at {position}: KL {kl} > {bound}"

check_deep_prompt

check_deep_prompt(model, rows: int, bound: float, chars: int = 1024, *, record)

Check the deepest prompt chunk (holding position max_seq_len - 1), at most rows rows of it, against the graph's own reference.

Within bound of its KL, every earlier position filled as check_deep_decode fills them.

Source code in iron/lm/testing.py

def check_deep_prompt(model, rows: int, bound: float, chars: int = 1024, *, record):
    """Check the deepest prompt chunk (holding position ``max_seq_len - 1``), at most ``rows`` rows of it, against the graph's own reference.

    Within ``bound`` of its KL, every earlier position filled as
    ``check_deep_decode`` fills them.
    """
    tokens = prompt(model, chars, 1)
    model.logits(tokens)
    C, L = model.config.prefill_chunk, model.config.max_seq_len
    deepest = (L - 1) // C
    rows = min(rows, L - deepest * C)
    chunk = np.resize(tokens, rows)
    ids = model.session.prompt_ids(chunk)
    begin = deepest * C
    values = dict(
        token=int(chunk[-1]),
        position=begin + rows - 1,
        start=begin,
        rows=rows,
        rope_start=model.rope_row(begin + rows, begin),
    )
    with _fill_caches(model, len(tokens), begin):
        start = time.perf_counter()
        got, _ = model(ids, **values)
        record("DeepPromptSeconds", time.perf_counter() - start)
        got = got.numpy()
        expected, _ = model.reference(ids, **values)
    kl, top1 = divergence(expected, got)
    record("DeepPromptKL", kl)
    record("DeepPromptTop1", top1)
    assert kl <= bound, f"the chunk at {begin}: KL {kl} > {bound}"

cli

main, a model's command line.

exports

exports(package: str, names)

Build a model package's __getattr__, each of names from its model module.

Imported on first use, so python -m <package>.model runs that module once.

Source code in iron/lm/cli.py

def exports(package: str, names):
    """Build a model package's ``__getattr__``, each of ``names`` from its model module.

    Imported on first use, so ``python -m <package>.model`` runs that
    module once.
    """

    def __getattr__(name):
        if name in names:
            return getattr(importlib.import_module(f"{package}.model"), name)
        raise AttributeError(f"module {package!r} has no attribute {name!r}")

    return __getattr__

sampling_arguments

sampling_arguments(parser: ArgumentParser) -> None

--temperature and --top-k, a Sampler's, at its defaults.

Source code in iron/lm/cli.py

def sampling_arguments(parser: argparse.ArgumentParser) -> None:
    """``--temperature`` and ``--top-k``, a ``Sampler``'s, at its defaults."""
    defaults = inspect.signature(Sampler).parameters
    parser.add_argument(
        "--temperature", type=float, default=defaults["temperature"].default
    )
    parser.add_argument("--top-k", type=int, default=defaults["top_k"].default)

parse_args

parse_args(parser: ArgumentParser, args=None) -> Namespace

parser's arguments with -v, which logs iron's INFO lines.

Such as each compile's.

Source code in iron/lm/cli.py

def parse_args(parser: argparse.ArgumentParser, args=None) -> argparse.Namespace:
    """``parser``'s arguments with ``-v``, which logs iron's INFO lines.

    Such as each compile's.
    """
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="log iron's INFO lines"
    )
    args = parser.parse_args(args)
    if args.verbose:
        logging.basicConfig()
        logging.getLogger("iron").setLevel(logging.INFO)
    return args

main

main(model_class: type[Model])

Run the command line of model_class, loaded as Model.load loads it.

Sample, or check it.

Source code in iron/lm/cli.py

def main(model_class: type[Model]):
    """Run the command line of ``model_class``, loaded as ``Model.load`` loads it.

    Sample, or check it.
    """
    parser = argparse.ArgumentParser(description=f"{model_class.__name__} on the NPU")
    parser.add_argument(
        "path",
        help="a directory holding config.json, the checkpoint and the "
        "tokenizer, or a Hugging Face repository's id",
    )
    for flags, kwargs in FLAGS:
        parser.add_argument(*flags, **kwargs)
    parser.add_argument("--prompt", metavar="TEXT", help="the text to continue")
    parser.add_argument(
        "--prompt-len",
        type=int,
        default=2048,
        help="without --prompt, characters of prompt.txt to prompt with, or "
        "with --random-weights a third as many random tokens (default: 2048)",
    )
    parser.add_argument(
        "--num-tokens", type=int, default=40, help="tokens to generate (default: 40)"
    )
    sampling_arguments(parser)
    check = parser.add_mutually_exclusive_group()
    check.add_argument(
        "--check-accuracy",
        action="store_true",
        help="instead of sampling, compare each step's logits with the CPU "
        "reference's, feeding both the reference's greedy token",
    )
    check.add_argument(
        "--check-determinism",
        type=int,
        metavar="ROUNDS",
        help="instead of sampling, run two prompts ROUNDS times each, "
        "alternating, and count the runs whose logits differ bitwise",
    )
    parser.add_argument(
        "--compare-host",
        action="store_true",
        help="after drawing on the device, generate again drawing on the host "
        "from the same seed and count the tokens that differ",
    )
    parser.add_argument(
        "--full-elf",
        action=argparse.BooleanOptionalAction,
        help="package each version as one full ELF, or with --no-full-elf "
        "dispatch every step on its own from one xclbin, as NPU1 does "
        "(default: a full ELF where the device runs one)",
    )
    args = parse_args(parser)
    if args.check_determinism is not None and args.check_determinism < 2:
        parser.error("--check-determinism compares rounds: give 2 or more")
    try:
        model = model_class.from_checkpoint(
            args.path, random_weights=args.random_weights, max_seq_len=args.max_seq_len
        )
    except ValueError as e:
        parser.error(str(e))
    if args.prompt is None:
        text, tokens = prompt_text(args.prompt_len), model.prompt(args.prompt_len)
    elif model.tokenizer is None:
        parser.error("--prompt takes the checkpoint's tokenizer: no --random-weights")
    else:
        text, tokens = args.prompt, model.tokenizer(args.prompt)
    if len(tokens) + args.num_tokens > model.config.max_seq_len:
        parser.error(
            f"a {len(tokens)}-token prompt and {args.num_tokens} more exceed "
            f"{model.config.max_seq_len} rows"
        )
    assert model.checkpoint is not None, "from_checkpoint gives it one"
    model.load_versions(release=model.checkpoint.release, full_elf=args.full_elf)

    if args.check_accuracy:
        results = accuracy(model, model.cpu_reference(), tokens, args.num_tokens)
        worst = int(np.argmax([kl for kl, _ in results]))
        for stat, value in kl_stats(results).items():
            where = f" (step {worst})" if stat == "Max" else ""
            print(f"[Accuracy] {stat} KL: {value:.6f}{where}")
        print(f"[Accuracy] Top-1 mismatches: {sum(not t for _, t in results)}")
    elif args.check_determinism is not None:
        # The second prompt is as much of the text as follows the first.
        prompts = [tokens, model.prompt(args.prompt_len, skip=args.prompt_len)]
        rounds = args.check_determinism
        n_differ = determinism(model, prompts, args.num_tokens, rounds)
        print(f"[Determinism] Differing runs: {n_differ}/{2 * (rounds - 1)}")
    else:

        def sampler():
            return Sampler(args.temperature, args.top_k)

        # The text, where there is a tokenizer to read it, else the ids.
        ids = model.tokenizer is None
        prompt = tokens if ids else text
        for piece in model.stream(prompt, args.num_tokens, sampler()):
            print(f" {piece}" if ids else piece, end="", flush=True)
        stats = model.stats
        assert stats is not None, "set once the last token is drawn"
        print(f"\n\n[Prefill] Time to first token: {stats.ttft:7.3f} s")
        if args.num_tokens > 1:
            print(f"[Decode]  Tokens per second:   {stats.tokens_per_s:7.3f}")
        if args.compare_host:
            drawn = model.generate(tokens, args.num_tokens, sampler())
            host = generate_on_host(model, tokens, args.num_tokens, sampler())
            differ = sum(a != b for a, b in zip(drawn, host))
            print(f"[Device] Tokens differing from the host loop: {differ}/{len(host)}")