Graph operators

An operation such as .softmax 1 has a constructor, .softmax, and a static attribute: axis 1. OpKind keeps both. OpTag keeps just the constructor, so code that needs an operation's name or parent count can ask for it without inventing an axis, shape, or convolution configuration.

Tensor parameters stay in the graph's external payload store. For example, .linear has no static attributes here, but its weight and bias still come from that store. The parent count describes dataflow: a linear node has one parent for its input tensor.

NN.IR.Graph adds node identities, dependency edges, and declared output shapes. The fixed torchlean.ir.v1 spelling of each constructor is defined separately in NN.Runtime.PyTorch.Wire.

structure NN.IR.HardMask : Type

A row-major Boolean mask carried by an IR operation.

The payload records its logical tensor shape separately from the flat array so graph validation can reject malformed serialized or programmatically constructed masks before evaluation. true means that the corresponding entry is allowed.

• shape : Spec.Shape
• allowed : Array Bool

@[instance_reducible]

instance NN.IR.instReprHardMask : Repr HardMask

def NN.IR.instReprHardMask.repr : HardMask → ℕ → Std.Format

def NN.IR.instDecidableEqHardMask.decEq (x✝ x✝¹ : HardMask) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.IR.instDecidableEqHardMask : DecidableEq HardMask

structure NN.IR.WindowConfig : Type

Per-axis geometry for pooling and convolution operators.

All three tensors describe the same spatial suffix of the input tensor. Keeping the geometry in one record prevents frontends from silently imposing a common stride or padding on every axis.

• spatialRank : ℕ

  Number of spatial axes.

• kernel : TorchLean.Tensor ℕ [self.spatialRank]

  Window extent along each spatial axis.

• stride : TorchLean.Tensor ℕ [self.spatialRank]

  Step along each spatial axis.

• padding : TorchLean.Tensor ℕ [self.spatialRank]

  Symmetric padding along each spatial axis.

@[instance_reducible]

instance NN.IR.instReprWindowConfig : Repr WindowConfig

def NN.IR.instReprWindowConfig.repr : WindowConfig → ℕ → Std.Format

def NN.IR.instDecidableEqWindowConfig.decEq (x✝ x✝¹ : WindowConfig) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.IR.instDecidableEqWindowConfig : DecidableEq WindowConfig

structure NN.IR.ConvConfigextends NN.IR.WindowConfig : Type

Shape metadata for an arbitrary-dimensional convolution.

Axes before channelAxis are preserved and mapped independently. The channel axis is replaced by outChannels; every following axis is spatial and is governed by window.

• spatialRank : ℕ
• kernel : TorchLean.Tensor ℕ [self.spatialRank]
• stride : TorchLean.Tensor ℕ [self.spatialRank]
• padding : TorchLean.Tensor ℕ [self.spatialRank]
• dilation : TorchLean.Tensor ℕ [self.spatialRank]

  Spacing between kernel elements along each spatial axis.

• paddingAfter : TorchLean.Tensor ℕ [self.spatialRank]

  Zero padding after the input along each spatial axis.

  The inherited padding field is the padding before the input. Keeping both sides explicit represents asymmetric padding without introducing rank-specific convolution variants.

• groups : ℕ

  Number of channel groups. 1 is an ordinary dense convolution.

• channelAxis : ℕ

  Axis containing the input channels.

• inChannels : ℕ

  Expected extent of the input-channel axis.

• outChannels : ℕ

  Extent of the output-channel axis.

@[instance_reducible]

instance NN.IR.instReprConvConfig : Repr ConvConfig

def NN.IR.instReprConvConfig.repr : ConvConfig → ℕ → Std.Format

def NN.IR.instDecidableEqConvConfig.decEq (x✝ x✝¹ : ConvConfig) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.IR.instDecidableEqConvConfig : DecidableEq ConvConfig

inductive NN.IR.OpKind : Type

Operation kinds in an op-tagged computation graph.

DecidableEq is derived (and therefore == is available) so passes can compare operation tags including their geometry payloads; the tensor-valued fields compare by row-major data.

• input : OpKind
• const (valueShape : Spec.Shape) : OpKind
• permute (perm : Array ℕ) : OpKind
• transpose (axis₁ axis₂ : ℕ) : OpKind
• detach : OpKind
• randUniform (seed : ℕ) : OpKind
• bernoulliMask (seed : ℕ) : OpKind
• add : OpKind
• sub : OpKind
• mulElem : OpKind
• abs : OpKind
• sqrt : OpKind
• inv : OpKind
• maxElem : OpKind
• minElem : OpKind
• maxPool (config : WindowConfig) : OpKind
• avgPool (config : WindowConfig) : OpKind
• broadcastTo (s₁ s₂ : Spec.Shape) : OpKind
• reduceSum (axis : ℕ) : OpKind
• reduceMean (axis : ℕ) : OpKind
• sum : OpKind
• matmul : OpKind
• linear : OpKind
• conv (config : ConvConfig) : OpKind
• batchNormEval (channelAxis channels : ℕ) : OpKind
• relu : OpKind
• tanh : OpKind
• sigmoid : OpKind
• exp : OpKind
• log : OpKind
• sin : OpKind
• cos : OpKind
• softplus : OpKind
• safeLog : OpKind
• softmax (axis : ℕ) : OpKind
• hardMaskedSoftmax (mask : HardMask) : OpKind
• layernorm (axis : ℕ) : OpKind
• reshape (inShape outShape : Spec.Shape) : OpKind
• flatten (s : Spec.Shape) : OpKind
• concat (axis : ℕ) : OpKind
• mseLoss : OpKind

@[instance_reducible]

instance NN.IR.instReprOpKind : Repr OpKind

def NN.IR.instReprOpKind.repr : OpKind → ℕ → Std.Format

@[instance_reducible]

instance NN.IR.instDecidableEqOpKind : DecidableEq OpKind

def NN.IR.instDecidableEqOpKind.decEq (x✝ x✝¹ : OpKind) : Decidable (x✝ = x✝¹)

structure NN.IR.ParentArity : Type

Permitted parent-count interval for an IR operation.

• min : ℕ

  Minimum number of parents required by the operation.

• max? : Option ℕ

  Maximum number of parents, or none when no finite upper bound is imposed.

@[instance_reducible]

instance NN.IR.instDecidableEqParentArity : DecidableEq ParentArity

def NN.IR.instDecidableEqParentArity.decEq (x✝ x✝¹ : ParentArity) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.IR.instReprParentArity : Repr ParentArity

def NN.IR.instReprParentArity.repr : ParentArity → ℕ → Std.Format

structure NN.IR.OpMetadata : Type

Structural metadata shared by all instances of an IR operation kind.

• tag : String

  Short tag used in diagnostics; artifact codecs define their own spelling.

• arity : ParentArity

  Permitted number of parent nodes.

def NN.IR.instDecidableEqOpMetadata.decEq (x✝ x✝¹ : OpMetadata) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.IR.instDecidableEqOpMetadata : DecidableEq OpMetadata

def NN.IR.instReprOpMetadata.repr : OpMetadata → ℕ → Std.Format

@[instance_reducible]

instance NN.IR.instReprOpMetadata : Repr OpMetadata

inductive NN.IR.OpTag : Type

Constructor identity of an NN.IR.OpKind, with the payload forgotten.

• input : OpTag
• const : OpTag
• permute : OpTag
• transpose : OpTag
• detach : OpTag
• randUniform : OpTag
• bernoulliMask : OpTag
• add : OpTag
• sub : OpTag
• mulElem : OpTag
• abs : OpTag
• sqrt : OpTag
• inv : OpTag
• maxElem : OpTag
• minElem : OpTag
• maxPool : OpTag
• avgPool : OpTag
• broadcastTo : OpTag
• reduceSum : OpTag
• reduceMean : OpTag
• sum : OpTag
• matmul : OpTag
• linear : OpTag
• conv : OpTag
• batchNormEval : OpTag
• relu : OpTag
• tanh : OpTag
• sigmoid : OpTag
• exp : OpTag
• log : OpTag
• sin : OpTag
• cos : OpTag
• softplus : OpTag
• safeLog : OpTag
• softmax : OpTag
• hardMaskedSoftmax : OpTag
• layernorm : OpTag
• reshape : OpTag
• flatten : OpTag
• concat : OpTag
• mseLoss : OpTag

@[instance_reducible]

instance NN.IR.instDecidableEqOpTag : DecidableEq OpTag

@[instance_reducible]

instance NN.IR.instReprOpTag : Repr OpTag

def NN.IR.instReprOpTag.repr : OpTag → ℕ → Std.Format

@[instance_reducible]

instance NN.IR.instInhabitedOpTag : Inhabited OpTag

def NN.IR.OpTag.all : List OpTag

Every semantic operator identity, in declaration order.

theorem NN.IR.OpTag.mem_all (tag : OpTag) : tag ∈ all

Every operator identity occurs in the enumeration.

def NN.IR.OpTag.toKind? : OpTag → Option OpKind

Construct an operation when its tag supplies all the static information.

For example, .relu becomes some .relu, while .softmax returns none because its axis is still missing. Tensor parameters such as linear weights are read separately from the payload store.

def NN.IR.OpTag.hasAttributes (tag : OpTag) : Bool

Whether constructing the operation needs an axis, shape, or other static attributes.

def NN.IR.OpTag.metadata : OpTag → OpMetadata

Structural metadata shared by every instance of an operation.

Parameters such as linear weights live outside the parent edges, so linear has arity one.

def NN.IR.OpKind.opTag : OpKind → OpTag

Forget the static attributes of an operation.

def NN.IR.OpKind.metadata (kind : OpKind) : OpMetadata

Structural metadata, obtained from the operation's constructor identity.

theorem NN.IR.OpKind.to_kind_op_tag (kind : OpKind) (h : kind.opTag.hasAttributes = false) : kind.opTag.toKind? = some kind

Forgetting attributes and reconstructing an attribute-free operation preserves it.

def NN.IR.OpKind.minParents (kind : OpKind) : ℕ

The minimum number of parent nodes expected by an OpKind.

def NN.IR.OpKind.maxParents? (kind : OpKind) : Option ℕ

An optional maximum number of parent nodes expected by an OpKind.

For concat, the verifier permits an arbitrary number of inputs (at least 2), so this returns none.

def NN.IR.OpKind.tag (kind : OpKind) : String

A short tag for error messages and debugging output.

def NN.IR.OpKind.describe : OpKind → String

Human-facing operation description including operation-local parameters.

tag is short and stable for grouping/log filtering. describe is for diagnostics: it prints axes, shapes, seeds, and convolution/pooling metadata so malformed graph dumps are useful without cross-referencing the original builder.