IR axis operations

IR axis-ops runtime tutorial.

This tutorial constructs and evaluates three IR operations with 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:])

Each operation accepts any in-bounds tensor dimension. The implementation may move that dimension to an innermost position while evaluating an optimized kernel, but its public semantics preserves the original shape and dimension numbering. Forward graph execution reports unsupported backend cases explicitly.

Run:

lake exe torchlean ir_axis_ops --execution eager

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

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

Tensor Shapes

These shapes illustrate an axis that is neither first nor last.

@[reducible, inline]

abbrev NN.Examples.DeepDives.IRAxisOps.baseTensorShape : 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

LayerNorm over axis 1 of a rank-three tensor, the interesting case because the normalized axis is neither the first nor the last.

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

Concatenation along axis 1, joining a [2, 3, 4] and a [2, 5, 4] tensor into [2, 8, 4].

Runner Helpers

def NN.Examples.DeepDives.IRAxisOps.firstScalars {α : Type} [TorchLean.Storage α] [ToString α] {σ : Spec.Shape} (tensor : TorchLean.Tensor α σ) : String

Print a compact preview of a tensor.

def NN.Examples.DeepDives.IRAxisOps.printEvaluation {α : Type} {σ : Spec.Shape} [TorchLean.Storage α] [Context α] [ToString α] (tag : String) (g : IR.Graph) (payload : IR.Payload α) (x : TorchLean.Tensor α σ) (outputId : Fin g.nodes.size) : IO Unit

Evaluate an example graph through the runtime API and print its selected output.

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

Run every axis-op example at scalar type α.

Generic in α so the tutorial can be run under native Float or under the bit-level IEEE model with no change to the graphs.

def NN.Examples.DeepDives.IRAxisOps.runSelected {α : Type} [TorchLean.Storage α] [Context α] [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

Entry point; --arithmetic picks the scalar type the whole tutorial runs at.