TorchLean API

NN.API.Neural.VisionLayers

Vision Layers #

This file provides named-field layer records for spatial operators. Tensors remain ordinary arbitrary-rank tensors; each operator states the trailing axes it consumes, while leading records any axes mapped pointwise by the layer.

Spatial layers #

Configuration shared by arbitrary-dimensional convolution layers.

  • outChannels :

    Number of output channels.

  • kernel : Vector d

    Kernel extent along each spatial axis.

  • stride : Vector d

    Step along each spatial axis.

  • padding : Vector d

    Symmetric zero-padding along each spatial axis.

  • kernelNonzero (i : Fin d) : self.kernel.get i 0

    Every kernel extent is positive.

  • strideNonzero (i : Fin d) : self.stride.get i 0

    Every stride is positive.

  • seedKernel :

    Seed for deterministic kernel initialization.

  • seedBias :

    Seed for deterministic bias initialization.

  • Initialization scheme for the kernel weights.

Instances For
    def TorchLean.nn.Internal.conv (leading : Spec.Shape := Spec.Shape.scalar) {d inChannels : } (spatial : Vector d) (cfg : Conv d) [NeZero inChannels] :
    Sequential (leading.concat (Spec.Shape.ofList (inChannels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (cfg.outChannels :: (Spec.convOutSpatial spatial cfg.kernel cfg.stride cfg.padding).toList)))

    Apply an arbitrary-dimensional convolution to the channel and spatial suffix of a tensor.

    The input suffix is (inChannels, spatial...). Any axes in leading are preserved; internally they are flattened into one runtime batch and restored after the convolution.

    Instances For

      Configuration shared by arbitrary-dimensional pooling layers.

      • kernel : Vector d

        Window extent along each spatial axis.

      • stride : Vector d

        Step along each spatial axis.

      • padding : Vector d

        Symmetric padding along each spatial axis.

      • kernelNonzero (i : Fin d) : self.kernel.get i 0

        Every window extent is positive.

      • strideNonzero (i : Fin d) : self.stride.get i 0

        Every stride is positive.

      Instances For
        def TorchLean.nn.Internal.maxPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (cfg : Pool d) :
        Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial cfg.kernel cfg.stride cfg.padding).toList)))

        Apply max pooling to the channel and spatial suffix of a tensor.

        Instances For
          def TorchLean.nn.Internal.avgPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (cfg : Pool d) :
          Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial cfg.kernel cfg.stride cfg.padding).toList)))

          Apply average pooling to the channel and spatial suffix of a tensor.

          Instances For
            def TorchLean.nn.Internal.globalAvgPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (spatialNonzero : ∀ (i : Fin d), spatial.get i 0) :
            Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.dim channels Spec.Shape.scalar))

            Global average pooling over every spatial axis, preserving the leading axes and channels.

            Instances For

              LayerNorm configuration for batched (batch x seqLen x embedDim) tensors.

              PyTorch analogue: torch.nn.LayerNorm. See https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html.

              • seedGamma :

                Seed for deterministic initialization of gamma (scale).

              • seedBeta :

                Seed for deterministic initialization of beta (shift).

              Instances For
                def TorchLean.nn.Internal.layerNormWith {batch seqLen embedDim : } (cfg : LayerNorm) (hSeq : seqLen > 0) (hEmbed : embedDim > 0) :

                Layer normalization over (batch × seqLen × embedDim) tensors, with explicit positivity proofs.

                This matches the common Transformer usage: normalize each token’s embedDim-vector independently, with learnable scale/shift parameters gamma and beta.

                PyTorch analogue: torch.nn.LayerNorm(embedDim) applied to a tensor of shape (batch, seqLen, embedDim).

                Call nn.layerNorm when NeZero can discharge the positivity proofs automatically.

                Instances For
                  def TorchLean.nn.Internal.layerNorm {batch seqLen embedDim : } (cfg : LayerNorm := { }) [NeZero seqLen] [NeZero embedDim] :

                  Layer normalization over (batch × seqLen × embedDim) tensors.

                  This normalizes each embedDim-vector (per batch element, per sequence position), and applies learned affine parameters gamma and beta.

                  PyTorch analogue: torch.nn.LayerNorm(embedDim) on a tensor shaped (batch, seqLen, embedDim).

                  Implementation note: TorchLean uses NeZero to ensure seqLen and embedDim are positive, avoiding degenerate shapes.

                  Instances For

                    RMSNorm configuration for batched (batch x seqLen x embedDim) tensors.

                    This is a common alternative to LayerNorm in modern transformer architectures.

                    • seedGamma :

                      Seed for deterministic initialization of gamma (scale).

                    Instances For
                      def TorchLean.nn.Internal.rmsNormWith {batch seqLen embedDim : } (cfg : RMSNorm) (hSeq : seqLen > 0) (hEmbed : embedDim > 0) :

                      RMS normalization over (batch × seqLen × embedDim) tensors, with explicit positivity proofs.

                      This is like LayerNorm but without mean subtraction: we scale by the root-mean-square over the embedDim axis, and apply a learned scale gamma.

                      PyTorch analogue: many libraries provide an RMSNorm(embedDim) module; conceptually it is applied to tensors shaped (batch, seqLen, embedDim).

                      Call nn.rmsNorm when NeZero can discharge the positivity proofs automatically.

                      Instances For
                        def TorchLean.nn.Internal.rmsNorm {batch seqLen embedDim : } (cfg : RMSNorm := { }) [NeZero seqLen] [NeZero embedDim] :

                        RMS normalization over (batch × seqLen × embedDim) tensors.

                        This normalizes by the root-mean-square over the embedDim axis (per batch element, per position), then applies a learned scale gamma.

                        Implementation note: TorchLean uses NeZero to ensure seqLen and embedDim are positive, avoiding degenerate shapes.

                        Instances For

                          Parameter initialization for affine channel normalization.

                          Instances For
                            def TorchLean.nn.Internal.Implementation.reshapeLayer (source target : Spec.Shape) (sameSize : source.size = target.size) :
                            LayerDef source target

                            A checked reshape layer used internally to flatten and restore spatial axes.

                            Instances For
                              def TorchLean.nn.Internal.Implementation.channelFirstKernel (leadingSize channels spatialSize : ) (kernel : LayerDef (Spec.Shape.dim leadingSize (Spec.Shape.dim channels (Spec.Shape.dim spatialSize (Spec.Shape.dim 1 Spec.Shape.scalar)))) (Spec.Shape.dim leadingSize (Spec.Shape.dim channels (Spec.Shape.dim spatialSize (Spec.Shape.dim 1 Spec.Shape.scalar))))) :
                              Sequential (Spec.Shape.dim leadingSize (Spec.Shape.dim channels (Spec.Shape.dim spatialSize Spec.Shape.scalar))) (Spec.Shape.dim leadingSize (Spec.Shape.dim channels (Spec.Shape.dim spatialSize Spec.Shape.scalar)))

                              Apply a channel-first normalization kernel after adding its singleton trailing axis.

                              Instances For

                                Flatten arbitrary spatial axes to the channel-first kernel representation.

                                Instances For

                                  Restore the original spatial axes after channel normalization.

                                  Instances For
                                    def TorchLean.nn.Internal.batchNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] :
                                    Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList)))

                                    Batch normalization over (leading..., channels, spatial...) for any spatial rank.

                                    Instances For
                                      def TorchLean.nn.Internal.instanceNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] :
                                      Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList)))

                                      Instance normalization over (leading..., channels, spatial...) for any spatial rank.

                                      Instances For
                                        def TorchLean.nn.Internal.groupNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : } (spatial : Vector d) (groups : ) (hGroups : groups > 0) (hGroupsLe : channels groups) (hDiv : channels % groups = 0) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] :
                                        Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList)))

                                        Group normalization over (leading..., channels, spatial...) for any spatial rank.

                                        Instances For

                                          Multi-head self-attention configuration.

                                          PyTorch analogue: torch.nn.MultiheadAttention (conceptually). See https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html.

                                          • numHeads :

                                            Number of attention heads.

                                          • headDim :

                                            Per-head embedding dimension.

                                          • seedW :

                                            Base seed for deterministic parameter initialization.

                                          • Projection-weight initialization. none retains Xavier-uniform initialization.

                                          • Optional initializer for the output projection.

                                            This is separate because deep residual stacks commonly scale the projection that writes back to the residual stream. When omitted, weightInit? is used.

                                          • outputBias : Bool

                                            Add a trainable bias after the output projection.

                                          Instances For

                                            Multi-head self-attention with an explicit nonzero sequence length proof.

                                            If mask is provided, it is a boolean attention mask of shape (n × n) (e.g. causal masking).

                                            Instances For

                                              Multi-head self-attention using NeZero to hide the nonzero sequence length proof.

                                              If mask is provided, it is a boolean attention mask of shape (n × n) (e.g. causal masking).

                                              Instances For