TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.API.Models.Mamba

Mamba Models #

Reusable configuration, model constructors, and text helpers for Mamba-style sequence models.

The trainable model path uses TorchLean autograd layers and therefore runs on the CPU and CUDA backends. The spec-backed deterministic helpers below are kept as small mathematical reference utilities; runnable training examples use the autograd constructor.

source
structure TorchLean.nn.models.MambaTextConfig :
Type

Configuration for byte-level Mamba-style language models.

Instances For
    source
    def TorchLean.nn.models.instReprMambaTextConfig.repr :
    MambaTextConfig → ℕ → Std.Format
    Instances For
      source
      @[implicit_reducible]
      instance TorchLean.nn.models.instReprMambaTextConfig :
      Repr MambaTextConfig
      source
      @[reducible, inline]
      abbrev TorchLean.nn.models.mambaTokenMat (cfg : MambaTextConfig) (seqLen : ℕ) :
      Spec.Shape

      Sequence-major one-hot token matrix shape.

      Instances For
        source
        @[reducible, inline]
        abbrev TorchLean.nn.models.mambaLogitMat (cfg : MambaTextConfig) (seqLen : ℕ) :
        Spec.Shape

        Output logits shape for byte-level causal language modeling.

        Instances For
          source
          def TorchLean.nn.models.mambaTextLm (cfg : MambaTextConfig) (seqLen : ℕ) :
          M (Sequential (mambaTokenMat cfg seqLen) (mambaLogitMat cfg seqLen))

          Trainable Mamba-style causal language model over one-hot token inputs.

          Architecture:

          mamba(seqLen, vocab, stateDim) → linear(stateDim → vocab) applied at every time step.

          The recurrent core is a gated diagonal state-space update implemented with autograd-covered TorchLean ops. Passing --device cuda to a runner that instantiates this model trains the same parameters on the CUDA backend.

          Instances For
            source
            def TorchLean.nn.models.mambaCenteredHash (seed modulus : ℕ) :
            Float

            Small deterministic initializer for spec-level reference blocks.

            Instances For
              source
              def TorchLean.nn.models.mambaVectorFloat {n : ℕ} (f : Fin n → Float) :
              Spec.Tensor Float (NN.Tensor.shapeOfDims [n])

              Build a vector tensor from an index function.

              Instances For
                source
                def TorchLean.nn.models.mambaMatrixFloat {m n : ℕ} (f : Fin m → Fin n → Float) :
                Spec.Tensor Float (NN.Tensor.shapeOfDims [m, n])

                Build a matrix tensor from an index function.

                Instances For
                  source
                  def TorchLean.nn.models.compactMambaFloat (cfg : MambaTextConfig) :
                  Models.MambaBlockSpec Float cfg.vocab cfg.stateDim cfg.vocab

                  Compact diagonal Mamba-style block for spec-level reference evaluation.

                  Instances For
                    source
                    def TorchLean.nn.models.selectiveMambaFloat (cfg : MambaTextConfig) :
                    Models.SelectiveMambaBlockSpec Float cfg.vocab cfg.stateDim cfg.ssmStateDim cfg.vocab cfg.convWidth

                    Full selective Mamba-style block with causal depthwise convolution and token-dependent scan parameters. This deterministic initializer is meant for reference evaluation rather than checkpoint-quality training.

                    Instances For
                      source
                      def TorchLean.nn.models.mambaTrainingOffsets (tokenCount seqLen windows : ℕ) :
                      List ℕ

                      Deterministic starting offsets for fixed-width language-model training windows.

                      Instances For