# 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.

```lisp
(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:

```lisp
(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

```lisp
(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

| argument      | contract                                                                                                                                                                 |
| ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| `body`        | a function whose first argument is the incoming state, followed by one argument per sequence; each is a scalar, or with `:state-axis` a vector                           |
| `initial`     | an `f32` tensor whose shape is the sequence shape with `:axis` removed; with `:state-axis`, its last axis is the state's width                                           |
| `sequences`   | a nonempty list or vector of `f32` tensors, all with exactly the same shape; with `:state-axis`, the same shape except the last axis, whose width is each sequence's own |
| `:axis`       | the sequence axis, default `0`; negative axes count from the end; its length is at least one                                                                             |
| `:state-axis` | the last axis of the sequences, `-1`; makes the state a vector. Absent, the state is a scalar                                                                            |
| `:reverse`    | `true` or `false`, default `false`; `true` runs the trips from the last position of `:axis` to the first                                                                 |

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`](https://sx.041.io/docs/explain.md#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.

---

SexpGPU documentation. Every page: https://sx.041.io/llms.txt