Iterator chains

An iterator chain computes something for every element along an axis at once, as pipe stages: n |> map(i => e) is a tensor whose first axis is n, with e computed for every index i along it. It’s the parallel counterpart of a for loop, which stays sequential: a loop’s iterations run one after another and may depend on each other, while the elements of an iterator stage are computed together and can’t. No Python loop runs.

dim D = 8
dim C = 4

fn scale_rows(x: [n D], w: [n]) -> [n D] = n |> map(i => x[i] * w[i])

fn per_row(x: [n D], w: [D C]) -> [n C] = x |> map(row => softmax(row @ w, over=C), over=n)

fn gram(x: [n D]) -> [n n] = x |> map(a => x |> map(b => sum(a * b), over=n), over=n)

fn table() -> [3 D] f32 = 3 |> map(i => D |> map(j => j to f32))

fn positive(x: [n D]) -> [] = x
    |> map(row => sum(row), over=n)
    |> filter(t => t > 0.0)
    |> sum

What a chain iterates

The input of the first iterator stage says what is iterated:

• A size: an axis name of the function’s types, such as n in x: [n D], a loop’s batch axis b, a dim, a positive whole number (3), or a sum or product of them (2 * n). Its elements are its indices, each a [] i64 Index[n]: a tensor, not a host int, since every index is computed at once. So x[i] takes x’s first axis at i, as indexing with any Index[n] does (Indexing), and x[i] on an axis of another size is E0507; i to f32 is its value as a float. The iterated axis is named like the size when the size is a name, and has no name otherwise.
• A range a..b, whose ends are sizes: the indices from a up to b, an axis of size b - a. 0..n is the same as n; a range that doesn’t start at 0 gives [] i64 elements, not Index.
• A tensor: its slices along an axis, each with that axis removed. The axis is the one named by over= on the stage; or, if the tensor has one axis, that one; or, if the tensor is the result of a map or filter, the axis it iterated, which comes first in it. Otherwise the stage needs over= (E0406), and an over= that isn’t one of its axes is E0405.

Anything else, such as an int value known only when the program runs, a [] tensor or a Python value, is E2401: the result’s axis is the one iterated, so its size must be known from the types. A Python value can be given a shape with as first.

Stages

zip’s other iteration is a size, a range or a tensor, iterated along its first axis if it’s the result of a map or filter, its only axis, or the axis with the iterated axis’s name (E0406 otherwise). Zipping pairs with a third iteration gives triples. Pairs, from zip or enumerate, must be taken apart by a map or filter whose lambda has a parameter for each (E3206); anything else, such as ending a chain with them or giving them to another stage, is E3205. filter of pairs keeps pairs.

The stages compose like any pipe: a map’s result is a tensor that keeps the iterated axis’s name, so a reduction can name it, as in x |> map(row => relu(row), over=n) |> sum(over=n). After a filter, whose axis has no name, a reduction without over= reduces it, with the others. Each pipe line gets its inlay hint: a pair shows as its two tensors, each with the iterated axis first.

Functions

The function of map and filter is one of:

• A lambda, v => body or (a, b) => body, with a parameter for each part of an element (one, or two for pairs). Each parameter is in scope only in the body, can’t have the name of an axis or a dim (E0303), and _ ignores it. The body is an expression, which may be another chain: 3 |> map(i => D |> map(j => j to f32)) is [3 D] f32, with i and j both in scope in the inner body. A lambda anywhere else is E3202: lambdas aren’t values.
• A pipe stage, applied to each element as |> would apply it (Pipes): a function’s name (map(gelu), map(normalize, over=n)), a call that the element is passed to first or as its _ (map(clamp(0.0, 1.0)), map(sub(1.0, _))), or to T (map(to f32)). These take one element, not pairs (E3206).

Independence

No element may depend on another, and computing one may not change anything; the checker enforces both:

• A lambda can’t read the result it builds. In y = n |> map(i => y[i - 1] + x[i]), where y isn’t bound before, y would be other elements of the result: E2402. Each element computed from the one before is a recurrence, which a for loop computes. If y is bound before the statement, the lambda reads that old value, which is fine.
• The function only computes a value (E2403). It can’t use mut arguments, Python (a call through import python, or an extern fn outside torch and math), print, save or load, a mode (dropout and other random draws, whose random numbers would be shared or differ by order), value_and_grad, optimizers, update, schedules, datasets, logs records or sinks, or a slice bound that depends on a parameter, since every element must have the same shape. Reading models, their fields, gradients and batches is fine. A map must give a tensor, not a model.
• Functions it calls, in the lambda or given by name, are checked the same way, at any depth, and also can’t change a model’s field, branch with an if on a tensor (which would read the device for each element) or loop over a Python value. The error is at the call, with the construct in the function as a secondary label. A function with mut parameters is only called with mut arguments, which are reported.

The input of a chain, zip’s other iteration and over= are computed once, before the stage, and aren’t checked this way.

How iterator chains run

Each map and filter is one batched computation, chosen by the checker:

• Plain tensor operations, when a lambda is elementwise in its parameters: made of the parameters, x[i] (when i is the index of a size), numbers, values that don’t use the parameters, arithmetic and comparisons (not @), to conversions of the dtype or domain, and elementwise functions (built-in, extern fn or user functions whose parameters have no shapes), with every part that uses a parameter shaped like the body or []. The body is computed once on whole tensors: a parameter becomes the whole iterated tensor, x[i] becomes x, and a [] part gets size-1 axes to line up. So n |> map(i => x[i] * w[i]) is x * w[:, None]. A function given by name that is elementwise, or to T without a shape, is applied to the whole tensor: x |> map(gelu, over=n) is gelu(x).
• torch.vmap otherwise: torch.vmap(lambda row: body)(x), over the iterated tensors with the iterated axis moved first, nested for a nested chain.

A size’s indices are torch.arange(n), and a range’s torch.arange(a, b). filter computes its mask the same way, then indexes with it. zip and enumerate only pair tensors up: enumerate’s indices are torch.arange of the axis’s length.

Either way autograd follows it, so gradients flow through a chain, value_and_grad included. Shape checks in a lambda (as [D], and the argument checks of functions it calls) check each element’s shape. A value check in a lambda checks every element’s values together, and fails as any value check does (Runtime checks).

Export. torch.export records torch.vmap, so an export fn can use iterator chains, and the axis they iterate can stay dynamic (Export). A filter gives an unknown size, which export handles as it does unknown sizes. In an exported program, a value check inside a lambda isn’t kept: torch.vmap can’t give every element’s values to one assertion. Check the result after the chain instead.

Padding. A map over a padded axis computes each position on its own, so it needs nothing when its function uses values with the padded axis only as the elements it’s given, or as x[i] with i an index of the axis. Any other use of the padded axis in a stage’s function, such as summing over it inside, and a filter, zip or enumerate over it, keeps the axis from being padded (H2001).

Several devices. In a function that runs on the workers, an iterator stage that uses a value split across the devices, along the batch axis or along H with split H, is E2404: it has no lowering that works on each device’s part. One that uses only whole values, such as a model’s params, is computed whole on every device. Rewrite one over the batch with operations on whole tensors, which the checker splits.

Loops that could be iterator chains

A sequential loop for i in 0..n:, over a whole axis n, gets hint H2401 when its iterations don’t depend on each other: every statement of its body binds a value of that iteration, or adds a term to an accumulator from before the loop (acc = acc + e or acc = e + acc); no term reads an accumulator; i is used only as an index, x[i]; and nothing in it does more than compute a value, as for a stage’s function. The help suggests acc = acc + (n |> map(i => e) |> sum(over=n)). A loop whose iterations read what the last one computed, that uses i as a number, or that prints, gets no hint.