IR axis operations

IR axis-ops runtime tutorial.

This example is a small regression guard for three ops where TorchLean’s IR uses an explicit axis:

• softmax axis (PyTorch: torch.softmax(x, dim=axis))
• concat axis (PyTorch: torch.cat(xs, dim=axis))
• layernorm axis PyTorch: F.layer_norm(x, normalized_shape=x.shape[axis:])

Why this tutorial exists:

• These three ops are easy to accidentally restrict to “last axis only” (because the spec primitives we reuse are last-axis).
• The denotational IR semantics supports the PyTorch meaning on any valid axis: it implements non-last axes by reshaping/permuting into a form the spec primitive already supports.
• The compiled IRExec backend is more conservative today. This tutorial runs compiled execution only for supported cases and prints an explicit skip for known backend gaps, instead of treating the backend covers more than it does.

Run:

lake exe torchlean ir_axis_ops --dtype float --backend eager

def NN.Examples.DeepDives.IRAxisOps.usage : String

Command-line help for the IR axis-ops tutorial.

Test Shapes

We keep shapes compact while still exercising the “axis is not last / not 0” code paths.

@[reducible, inline]

abbrev NN.Examples.DeepDives.IRAxisOps.baseRankThreeShape : Spec.Shape

@[reducible, inline]

abbrev NN.Examples.DeepDives.IRAxisOps.widerMiddleAxisShape : Spec.Shape

@[reducible, inline]

abbrev NN.Examples.DeepDives.IRAxisOps.concatenatedMiddleAxisShape : Spec.Shape

Small IR Graphs

def NN.Examples.DeepDives.IRAxisOps.softmaxMiddleAxisGraph : IR.Graph

def NN.Examples.DeepDives.IRAxisOps.layerNormMiddleAxisGraph : IR.Graph

def NN.Examples.DeepDives.IRAxisOps.concatMiddleAxisGraph : IR.Graph

Runner Helpers

def NN.Examples.DeepDives.IRAxisOps.checkIR (tag : String) (g : IR.Graph) : IO Unit

def NN.Examples.DeepDives.IRAxisOps.firstScalars {α : Type} [ToString α] : List α → String

def NN.Examples.DeepDives.IRAxisOps.runOne {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] [TorchLean.Runtime.FromFloat α] (tag : String) (g : IR.Graph) (payload : IR.Payload α) (inputShape : Spec.Shape) (x : Spec.Tensor α inputShape) (outputId : Fin g.nodes.size) (forwardArtifact : Bool := true) : IO Unit

def NN.Examples.DeepDives.IRAxisOps.runOnce {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] [TorchLean.Runtime.FromFloat α] : IO Unit

def NN.Examples.DeepDives.IRAxisOps.runSelected {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] [TorchLean.Runtime.FromFloat α] (rest : List String) : IO Unit

Runtime-selected entrypoint body for the axis-ops tutorial.

def NN.Examples.DeepDives.IRAxisOps.main (args : List String) : IO Unit