Batch mapping

A backend-independent recursion for applying an operation pointwise over the outer axis of a shape-typed reference. Eager and typed graph execution supply their own slicing, reshaping, and concatenation operations; the traversal itself remains shared.

@[irreducible]

def Runtime.Autograd.mapBatch {m : Type → Type} [Monad m] {Ref : Spec.Shape → Type} {σ τ : Spec.Shape} (empty : m (Ref (Spec.Shape.dim 0 τ))) (slice : {n : ℕ} → Ref (Spec.Shape.dim n σ) → (start len : ℕ) → start + len ≤ n → m (Ref (Spec.Shape.dim len σ))) (reshape : {s₁ s₂ : Spec.Shape} → Ref s₁ → s₁.size = s₂.size → m (Ref s₂)) (concat : {n k : ℕ} → Ref (Spec.Shape.dim n τ) → Ref (Spec.Shape.dim k τ) → m (Ref (Spec.Shape.dim (n + k) τ))) (f : Ref σ → m (Ref τ)) {n : ℕ} (x : Ref (Spec.Shape.dim n σ)) : m (Ref (Spec.Shape.dim n τ))

Apply a reference-level operation independently to every entry of a leading axis.

The callbacks isolate the four structural operations needed by the recursion. Device backends may replace this reference traversal with a fused primitive when the fused operation has the same per-entry semantics.

The axis is split in half and each half is mapped recursively, so f still runs on entries 0, ..., n - 1 in order while every entry is sliced and concatenated once per level of a balanced tree. Peeling one entry at a time would copy the remaining tail at every step.