IR axis operations

Evaluate axis operations with the IR semantics.

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:])

The axis is part of each operation's meaning. The evaluator checks graph and input shapes, then computes the selected output with the IR semantics, in Lean tensor operations. This example does not dispatch to LibTorch or select a GPU device; --arithmetic chooses the scalar implementation.

Run:

scripts/lake.sh exe torchlean ir_axis_ops --arithmetic ieee

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

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

Small IR Graphs

def NN.Examples.DeepDives.IRAxisOps.softmax (axis : ℕ := 1) : IR.Graph

Softmax along a chosen axis of the example's [2, 3, 4] input.

def NN.Examples.DeepDives.IRAxisOps.layerNorm (axis : ℕ := 1) : IR.Graph

LayerNorm over the suffix starting at axis. The default normalizes [3, 4] together, rather than just the final dimension.

def NN.Examples.DeepDives.IRAxisOps.concat (axis : ℕ := 1) : IR.Graph

Join [2, 3, 4] and [2, 5, 4] along axis one. Other axis choices demonstrate shape rejection: the remaining dimensions must agree and the declared output must match the inferred shape.

Runner Helpers

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

Print a compact preview of a tensor.

def NN.Examples.DeepDives.IRAxisOps.print {α : 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 with Lean tensor operations and print its selected output.

def NN.Examples.DeepDives.IRAxisOps.run {α : 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.main (args : List String) : IO Unit

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