Leading-Axis Reference Maps

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

@[irreducible]

def Runtime.Autograd.mapLeadingAxisWith {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 : ℕ) → len + start ≤ 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.