Nonlinear element scans

element-scan runs a nonlinear recurrence independently at every position outside one sequence axis. The body is traced once; the sequence length is static and the device loops over it at runtime.

(element-scan (lambda (h x a) (tanh (+ x (* a h))))
              initial
              (list x (broadcast-to decay (shape x)))
              :axis 1)

For x of shape [batch time width], initial has shape [batch width] and the result [batch time width]. At every batch and width position this computes h[t] = tanh(x[t] + decay[t] * h[t-1]) from h[-1] = initial; each result position holds the state after that trip. (slice states 1 (- time 1) time) selects the final state with its time axis kept. State lasts for this call only.

With :state-axis, the state of each lane is a vector along the last axis, and the body reads and returns whole vectors:

(element-scan (lambda (h x a)
                (let ((y (+ (* a h) x)))
                  (/ y (maximum 1.0 (sqrt (sum (* y y) :axes [-1]))))))
              (zeros [batch heads 16])
              (list x (broadcast-to decay [batch time heads 1]))
              :axis 1
              :state-axis -1)

For x of shape [batch time heads 16], each of the batch * heads lanes carries a state of 16 values, and every trip caps the norm of the whole state at one: h[t] = y / max(1, |y|) with y = decay[t] * h[t-1] + x[t]. The result has shape [batch time heads 16].

With :reverse true, the trips run from the last position of the axis to the first: h[t] = f(h[t+1], x[t]) from h[time] = initial, and each result position still holds the state after its own trip, so the final state is at position 0, (slice states 1 0 1). A right-to-left max-plus sweep is

(element-scan (lambda (h x c) (maximum (- h c) x))
              initial
              (list x cost)
              :axis 1
              :reverse true)

It computes the same values as flipping the sequences along :axis, scanning forward and flipping the result back, without the two flips.

Arguments and body

The body returns one f32 value of the state's shape. It accepts pointwise arithmetic, transcendental functions, comparisons, where, scalar constants, boolean intermediates and casts between f32 and bool; scalar reshape and broadcast are identities. sexpgpu ir lists the body's operations under the element_scan node.

With :state-axis, the body's values are scalars and vectors of 1 to 16 values: the state, each sequence's values at the lane and trip, and anything broadcast from a scalar or a one-value vector. sum and max reduce a vector to a scalar, or with :keepdims true to a one-value vector that broadcasts back against the state. Arithmetic between a scalar and a vector broadcasts as elsewhere, so a sequence of width one is a per-trip gate for the whole state.

Parameters and context values enter through the sequence list; broadcast a shared parameter to the sequence shape before the call. A closure may capture compile-time numbers and functions. Reading an outer tensor implicitly is E-FLOW-002; reading a parameter or a context scalar implicitly is E-SCAN-001.

The body cannot multiply matrices, index another lane, nest a scan, stop a gradient, draw random values, or report metrics, diagnostics, guards or counters; do those before or after the scan. A scalar state's body cannot reduce; a vector state's body reduces only its vectors. There is one state per lane and one collected output per trip: no early exit or dynamic shape.

Precision and gradients

The recurrence always computes in f32, including in a :bf16 run or region; values computed outside it keep their own rounding, so a broadcast parameter may already be rounded by the outer policy. with-numerics :precision :bf16 inside the body is rejected.

Gradients flow to initial and every sequence. A broadcast parameter's gradient sums its contributions across time and the other broadcast axes; repeated reads and repeated sequence arguments accumulate as elsewhere. A vector state's gradient is another element scan, run from the last trip back to the first, that applies the body's vector-Jacobian product once a trip without forming the trip's Jacobian. A :reverse scan's gradient runs from the first trip to the last. The gradient is built from ordinary tensor operations and element scans, so smooth bodies support further gradient calls. Comparisons and boolean casts have no derivative, where follows the selected branch, and the usual tie rules apply. Nothing defines derivatives at branch boundaries or removes the numerical limits of higher derivatives through long recurrences.

Storage and execution

For T trips, L independent lanes, a state of n values and B body operations:

• Forward does O(T * L * B * n) work and stores T * L * n states.
• Backward keeps the state sequence and needs O(T * L * B * n) intermediate storage, recomputing the body over the saved states. There is no checkpoint option or recomputation interval. A final-state slice still keeps the whole state sequence for its backward pass.
• The graph size depends on the body, not on T.
• Lanes run in parallel and trips in sequence, so a long recurrence with few lanes is serial on the device. A :reverse scan costs the same as a forward one.
• Metal runs scalar states only, in either direction; it refuses a vector state when the graph is lowered, with a message naming CUDA and CPU.
• On CPU each trip of each lane is a separate interpreter call, far slower than the same number of pointwise operations; train long or wide scans on Metal or CUDA.