Rounded evaluation-mode BatchNorm

The real operation broadcasts stored channel statistics across the spatial shape and applies Spec.batchNormInference independently to every leading sample. The flat-index lemmas below connect that operation to the channel selected by the executable transfer.

def NN.MLTheory.CROWN.Graph.DirectedBackward.normalizationChannel (leading : Spec.Shape) {channels : ℕ} {spatial : Spec.Shape} : (leading.concat (Spec.Shape.dim channels spatial)).Coord → Fin channels

The channel coordinate after an arbitrary leading batch shape.

theorem NN.MLTheory.CROWN.Graph.DirectedBackward.axisCoordinateOfFlat_normalizationChannel (leading spatial : Spec.Shape) (channels : ℕ) (c : (leading.concat (Spec.Shape.dim channels spatial)).Coord) : axisCoordinateOfFlat (leading.concat (Spec.Shape.dim channels spatial)) leading.rank ↑c.linearize % channels = ↑(normalizationChannel leading c)

The executable division/modulo channel index is the actual tensor channel coordinate.

theorem NN.MLTheory.CROWN.Graph.DirectedBackward.normalization_shape_decomposition {s : Spec.Shape} {axis channels : ℕ} (h : s.toList[axis]? = some channels) : ∃ (leading : Spec.Shape) (spatial : Spec.Shape), axis = leading.rank ∧ s = leading.concat (Spec.Shape.dim channels spatial)

A present channel axis decomposes its shape into leading, channel, and spatial factors.

theorem NN.MLTheory.CROWN.Graph.DirectedBackward.batchNormInference_normalization_apply {channels : ℕ} {spatial : Spec.Shape} (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels spatial)) (mean variance gamma beta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ) (i : Fin channels) (c : spatial.Coord) : TorchLean.Tensor.Internal.Rep.get (Spec.batchNormInference x mean variance gamma beta epsilon) (i, c) = (TorchLean.Tensor.Internal.Rep.get x (i, c) - mean.getScalar i) / √(max (variance.getScalar i) 0 + epsilon) * gamma.getScalar i + beta.getScalar i

The coordinate formula for the actual inference specification, including variance clamping.

noncomputable def NN.MLTheory.CROWN.Graph.DirectedBackward.batchNormRealScalar {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] [LawfulBoundOps α] (config : IR.BatchNormEvalParams α) (channel : Fin config.c) (x : ℝ) : ℝ

Read the actual fixed running statistics and affine scalars for one channel.

noncomputable def NN.MLTheory.CROWN.Graph.DirectedBackward.batchNormRealValue {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] [LawfulBoundOps α] (leading spatial : Spec.Shape) (config : IR.BatchNormEvalParams α) (f : ℕ → ℝ) : TorchLean.Tensor ℝ (leading.concat (Spec.Shape.dim config.c spatial))

The real, channel-wise Spec operation applied across the same leading axes as IR evaluation.

def NN.MLTheory.CROWN.Graph.DirectedBackward.BatchNormRealEquation {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] [LawfulBoundOps α] (s : Spec.Shape) (axis : ℕ) (config : IR.BatchNormEvalParams α) (f g : ℕ → ℝ) : Prop

The mathematical BatchNorm node equation does not mention interval boxes.

theorem NN.MLTheory.CROWN.Graph.DirectedBackward.BatchNormRealEquation.apply {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] [LawfulBoundOps α] {s : Spec.Shape} {axis : ℕ} {config : IR.BatchNormEvalParams α} {f g : ℕ → ℝ} (heq : BatchNormRealEquation s axis config f g) (hshape : s.toList[axis]? = some config.c) (k : Fin s.size) (hc : axisCoordinateOfFlat s axis ↑k % config.c < config.c) : g ↑k = batchNormRealScalar config ⟨axisCoordinateOfFlat s axis ↑k % config.c, hc⟩ (f ↑k)

The Spec equation yields the scalar expression at the engine's selected flat channel.

theorem NN.MLTheory.CROWN.Graph.DirectedBackward.ibpBatchNormEval?_encloses {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] [LawfulBoundOps α] [NonlinearBoundOps α] [LawfulNonlinearBoundOps α] {s : Spec.Shape} {axis : ℕ} {config : IR.BatchNormEvalParams α} {input output : FlatBox α} {f g : ℕ → ℝ} (hx : RowEncloses input s.size f) (heq : BatchNormRealEquation s axis config f g) (hout : ibpBatchNormEval? s axis config input = some output) : RowEncloses output s.size g

Every successful evaluation-mode BatchNorm transfer encloses the real Spec operation with the stored running statistics, scale, bias, and epsilon on every accepted channel axis.