sexpgpu/axes

The standard named-axis module, built in: (require "sexpgpu/axes" rearrange broadcast-axes). Both macros resolve names and dimensions at compile time and emit plain reshape, transpose and, for new axes, broadcast-to, leaving out identities (including a broadcast whose new axes all have size 1). The input expression runs once, then each size expression once in declaration order.

(require "sexpgpu/axes" rearrange broadcast-axes)

; [batch time width] -> [batch heads time depth], width = heads * depth.
(rearrange x
           [batch time (heads depth)]
           [batch heads time depth]
           :sizes ((heads 8)))

; Merge the heads back, preserving their order within width.
(rearrange attention [batch heads time depth] [batch time (heads depth)])

; Repeat a [time depth] table for every batch and head.
(broadcast-axes table
                [time depth]
                [batch heads time depth]
                :sizes ((batch batch-size) (heads 8)))

Patterns

A pattern is a literal vector. Each entry is an axis symbol or a nonempty flat list of axis symbols. A list groups consecutive axes into one tensor dimension. The rightmost axis in a group varies fastest. Each axis occurs once per pattern. The source pattern has exactly one entry per input dimension.

Labels match by their symbol names across modules. A wrapper macro can supply one pattern while its caller supplies the other pattern or size declarations. Identically named labels count as duplicates even when different modules created their symbols. Input and size expressions retain their lexical bindings.

rearrange requires the same named axes in both patterns. It can split, merge, and permute axes. broadcast-axes also permits new target axes, each with an explicit size. Neither macro drops axes or performs reductions. Existing axes keep their sizes, including singleton axes.

[] describes a scalar. A singleton dimension has a name with size 1, such as [unit] with :sizes ((unit 1)). Anonymous axes, numeric axes, _, ellipses, nested groups, and dynamically computed patterns are unsupported.

Sizes and inference

:sizes is a literal list of (axis expression) pairs. Expressions can refer to caller bindings, including knobs. Each value must be a positive integer. An integral float such as 2.0 is accepted and converted to an integer. Duplicate entries and entries for axes absent from the target are errors.

A source group can omit one size. The compiler divides that tensor dimension by the product of the declared sizes, after checking exact divisibility. A standalone source axis therefore needs no declaration. If all sizes in a source group are declared, their product must equal the tensor dimension.

For example, a dimension of 24 matched by (heads depth) with heads = 8 infers depth = 3. A dimension of 25 fails. Supplying neither size is ambiguous and fails even if the target repeats the same group. Each group resolves independently. Target grouping supplies no extra size constraints.

Diagnostics

Invalid layouts fail during compilation with the macro call as the primary source location. The diagnostic includes the offending axis, group, or size. For an ambiguous split, the message names the missing sizes:

axes: ambiguous source group (heads depth); supply :sizes for all but one axis, missing (heads depth)

For a nondivisible split, it includes the tensor dimension and known product:

axes: source group (heads depth) dimension 25 must be positive and divisible by 8