Expand description
Give a model’s dataflow a batch axis, on axis 0 of every input and output, so that a turn can seat several callers in one run.
Every op the axis reaches is asked to host it, through the same
TypedOp::change_axes protocol declutter propagates axis changes with, and
the answer is accepted only if the op carries the axis through to all of
its outputs. That is a stricter contract than the protocol’s own, which only
owes semantics at extent 1: an EinSum offered an extra leading axis on an
input takes it as a fresh contraction label, summing an axis of extent one
away, which is exact there and sums the seats together once the extent is the
batch. So the batched EinSum is built here instead, with the label on the
output as well, and an op that answers with anything but the axis carried
through is refused by name.
An input the caller lists as shared is one every seat reads whole – a
table, a language id. It keeps its own shape, and the ops that need their
operands to agree on rank read it through an AxisOp::Add(0), so it
broadcasts across the turn.
An input already sized by the batch symbol has that axis moved to the
front, as an interface change plus the inverse AxisOp::Move on the wire
behind it: an export putting the batch inside ([T, B, H] being the common
shape) becomes a graph edit declutter can absorb rather than a transpose the
runtime would pay per turn. Only axis 0 makes a seat’s values a contiguous
run, which is what the laned runtime addresses and what stacking a turn’s
answer back relies on.
What this does not touch is state: a seat’s own history is the laned runtime’s lane addressing, and a batch axis on the dataflow does not imply one. Batchify gives you seats, lane addressing gives each seat its own past.
Structs§
Functions§
- batchify
- Batch axis on axis 0 of every input but the
sharedones, sized bybatch.