TorchLean API

NN.Runtime.Autograd.Model.Norm

Norm #

Normalization programs built from Ops, so they can run eagerly or be recorded in a typed graph. Spatial dimensions are represented by an arbitrary Shape; the same definitions cover vectors, images, volumes, and higher-rank data.

BatchNorm has separate training and evaluation programs here. Training computes statistics from the input, while evaluation receives the stored mean and variance as arguments. Layers.batchNorm owns those running-statistics buffers and updates them through Layer.updateBuffers.

Flatten an arbitrary spatial shape while preserving the batch and channel axes.

Flatten an arbitrary spatial shape while preserving the channel axis.

def Runtime.Autograd.Model.Norm.Internal.broadcastChannelToBatchSpatial {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] (batch channels spatialSize : ) (x : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) :
m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels (Spec.Shape.dim spatialSize Spec.Shape.scalar))))

Repeat a channel vector over the batch and flattened spatial axes.

Instances For
    def Runtime.Autograd.Model.Norm.rmsNorm {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {leading : Spec.Shape} {width : } (hWidth : width > 0) (x : RefTy m α (leading.appendDim width)) (gamma : RefTy m α (Spec.Shape.dim width Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
    m (RefTy m α (leading.appendDim width))

    Normalize each final-axis vector by its root mean square, then apply gamma.

    For each leading index, we compute meanSq = mean(x * x) and divide its vector by sqrt(max(meanSq, 0) + ε). The scale gamma has shape [width] and is broadcast across the leading axes. We square the entries of x directly, without first subtracting their mean.

    width must be positive. If a leading axis is empty, the result is an empty tensor of the same shape, so we return it before constructing the reduction.

    Instances For
      def Runtime.Autograd.Model.Norm.l2Normalize {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {leading : Spec.Shape} {width : } (hWidth : width > 0) (x : RefTy m α (leading.appendDim width)) (epsilon : RefTy m α Spec.Shape.scalar) :
      m (RefTy m α (leading.appendDim width))

      Divide each final-axis vector by sqrt(sum(x * x) + epsilon).

      epsilon is a scalar tensor reference, shared by all vectors. It is added to the squared norm before the square root, so its effect follows this formula even for vectors whose norm is very small. An empty leading axis produces an empty output of the same shape.

      Instances For
        def Runtime.Autograd.Model.Norm.instanceNorm {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {batch channels : } {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (x : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) (gamma beta : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
        m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)))

        Normalize each sample and channel using its own spatial mean and variance.

        We flatten the spatial axes to a vector of length spatial.size, subtract that vector's mean, and divide by sqrt(max(mean((x - mean) * (x - mean)), 0) + ε). The batch and channel axes remain separate throughout this calculation.

        gamma and beta each have one entry per channel. They are broadcast across the batch and spatial positions before restoring the original shape. All statistics come from the current input.

        Instances For
          def Runtime.Autograd.Model.Norm.groupNorm {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {batch channels groups : } {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (hGroups : groups > 0) (hGroupsLe : channels groups) (hDiv : channels % groups = 0) (x : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) (gamma beta : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
          m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)))

          Normalize equal, contiguous channel groups independently within each sample.

          The input is viewed as [batch, groups, channelsPerGroup * spatial.size]. Reducing the last axis therefore combines the channels and spatial positions belonging to one group. We subtract the group mean and divide by sqrt(max(groupVariance, 0) + ε), using the mean squared deviation for groupVariance.

          After normalization, we restore the channel axis and apply gamma and beta, which each have one entry per channel. The shape hypotheses ensure that every group has the same positive size.

          Instances For
            def Runtime.Autograd.Model.Norm.batchNormTrainStats {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {batch channels : } {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (x : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) (gamma beta : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
            m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)) × RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar) × RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar))

            Normalize a batch and return (output, mean, variance).

            Statistics have shape [channels]. We first average over spatial positions and then over the batch, so every entry of a channel contributes equally. The variance is the mean squared deviation over batch * spatial.size entries, clamped below by zero.

            The output uses sqrt(variance + ε), followed by the per-channel scale gamma and bias beta. The returned variance is the biased estimate used in this forward pass; a running-statistics update can apply Layers.unbiasedRunningVariance before storing it.

            Instances For
              def Runtime.Autograd.Model.Norm.batchNormTrain {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {batch channels : } {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (x : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) (gamma beta : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
              m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)))

              Batch normalization using the current batch and spatial statistics.

              This returns the output of batchNormTrainStats. Use that function when the caller also needs the per-channel mean and biased variance, for example to update running statistics.

              Instances For
                def Runtime.Autograd.Model.Norm.batchNormEval {α : Type} [TorchLean.Storage α] [Context α] {m : TypeType} [Monad m] [Ops m α] {batch channels : } {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (x : RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) (gamma beta mean var : RefTy m α (Spec.Shape.dim channels Spec.Shape.scalar)) (ε : α := TorchLean.normalizationEpsilon) :
                m (RefTy m α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)))

                Normalize with supplied per-channel mean and variance.

                For each channel, the output is gamma * (x - mean) / sqrt(max(var, 0) + ε) + beta, broadcast over the batch and spatial axes. This program reads the supplied references and leaves buffer updates to the caller. Layers.batchNorm uses it in evaluation mode with the stored running statistics.

                Instances For