24: Iterator chains: what they iterate and their independence (comprehensions before CTL-3)

E2401

Iteration over something that isn’t an axis or a tensor. An iterator chain starts with a size (an axis name, a dim, a positive number, or a sum or product of them), a range whose ends are sizes, or a tensor with at least one axis, because the result’s axis is the one iterated and its size must be known from the types. An int known only when the program runs, a [] tensor or a Python value isn’t one. (Until CTL-3 this was the code of a comprehension over something other than an axis.)

dim D = 4

fn firsts(x: [n D], k: int) -> [n D] = k |> map(i => x[0])

Iterate the axis, n |> map(i => x[i]), and slice the result to keep part of it; or use a for loop to repeat something a run-time number of times. See What a chain iterates.

E2402

Iterator stage that reads the result it builds. An iterator stage computes every element at once, so no element can read another; a lambda that uses the name the chain is being bound to, when nothing had that name before, would read other elements of the result.

dim D = 4

fn running(x: [n D]) -> [n D]:
    y = n |> map(i => y[i] * 0.5 + x[i])
    return y

Each element computed from the one before is a recurrence: compute it with a for loop, which runs in order. See Independence.

E2403

Iterator stage whose function does more than compute a value. The function of a map or filter runs for every element at once, so it may only compute a value: no mut arguments, Python calls, print, checkpoints, modes (random draws), value_and_grad, optimizers, schedules, datasets, logs or sinks, no slice bounds that depend on a parameter, and no function that changes a model, branches on a tensor or loops over a Python value. A map’s function must give a tensor.

dim D = 4

fn noisy(x: [n D], mode: Mode) -> [n D] = x |> map(row => dropout(row, 0.1, mode), over=n)

Do it before or after the chain: here, dropout(x, 0.1, mode) on the whole tensor. See Independence.

E2404

Iterator stage that uses values split across devices. With several devices, a value with the batch axis, or with split H a value with an axis of size H, is split across them, each holding a part. An iterator stage in a function on the workers has no lowering that works on the parts, so it can’t use such a value.

devices cuda[0 1]

dim D = 4

data Ex:
    x: [D]

fn doubled(batch: [b] Ex) -> [b D] = batch.x |> map(row => row * 2.0, over=b)

Compute it with operations on whole tensors, such as batch.x * 2.0, which the checker splits, or declare one device. See How iterator chains run.

H2401

Sequential loop whose iterations are independent. A for i in 0..n: over a whole axis whose iterations only add terms that don’t depend on each other to accumulators, as in acc = acc + f(x[i]), could compute every term at once. The loop is correct as it is: this is a hint.

dim D = 4

fn total(x: [n D]) -> [D]:
    acc = x[0] * 0.0
    for i in 0..n:
        acc = acc + relu(x[i])
    return acc

hint[H2401]: the iterations of this loop over `n` don't depend on each other

Compute the terms with map and add them up: acc = acc + (n |> map(i => relu(x[i])) |> sum(over=n)). See Loops that could be iterator chains.