TorchLean API

Docs HomeInstallationGuideExamplesGraphsTools

NN.Runtime.Autograd.Model.Layers.Attention

TorchLean NN: Attention #

source
def Runtime.Autograd.Model.Layers.Attention.forward {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch n modelWidth heads headWidth : ℕ} (queryWeight keyWeight valueWeight : Ref [modelWidth, heads * headWidth]) (outputWeight : Ref [heads * headWidth, modelWidth]) (queryBias keyBias valueBias : Option (Ref [heads * headWidth])) (outputBias : Option (Ref [modelWidth])) (probability : Option (Ref Spec.Shape.scalar)) (input : Ref [batch, n, modelWidth]) (mask : Option (TorchLean.Tensor Bool [n, n])) (seed : ℕ) (training : Bool) :
m (Ref [batch, n, modelWidth])

Attention with optional affine projection biases and dropout on softmax probabilities.

The operation order is projection, head splitting, scaled dot products, masked softmax, probability dropout, value aggregation, and output projection. Dropout therefore removes individual query/key contributions before values are mixed.

Instances For
    source
    def Runtime.Autograd.Model.Layers.Attention.optionalShapes (enabled : Bool) (shape : Spec.Shape) :
    List Spec.Shape

    Include one state slot only when the corresponding option is enabled.

    Instances For
      source
      def Runtime.Autograd.Model.Layers.Attention.optionalState (enabled : Bool) {shape : Spec.Shape} (value : TorchLean.Tensor Float shape) :
      TorchLean.TensorPack Float (optionalShapes enabled shape)

      Initial value for an optional state slot, with no placeholder when it is disabled.

      Instances For
        source
        def Runtime.Autograd.Model.Layers.Attention.optionalInit (enabled : Bool) (shape : Spec.Shape) (init : Module.RuntimeInit.FloatInit) :
        Module.RuntimeInit.Plan (optionalShapes enabled shape)

        Native initialization follows the same optional shape list as the ordinary state pack.

        Instances For
          source
          def Runtime.Autograd.Model.Layers.Attention.takeOptional {Ref : Spec.Shape → Type} (enabled : Bool) (shape : Spec.Shape) {rest : List Spec.Shape} :
          RefList Ref (optionalShapes enabled shape ++ rest) → Option (Ref shape) × RefList Ref rest

          Read one optional state slot and return the remaining, still shape-indexed references.

          Instances For
            source
            def Runtime.Autograd.Model.Layers.attention (batchSize sequenceLength modelWidth headCount headWidth : ℕ) {sequenceLengthNonzero : sequenceLength ≠ 0} (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed : ℕ := 0) (weightInitialization? outputWeightInitialization? : Option Torch.Init.Scheme := none) (mask : Option (TorchLean.Tensor Bool [sequenceLength, sequenceLength]) := none) (inputBias outputBias : Bool := false) (dropout? : Option Float := none) (dropoutSeed : ℕ := 0) :
            Layer [batchSize, sequenceLength, modelWidth] [batchSize, sequenceLength, modelWidth]

            Configurable affine self-attention with optional probability dropout.

            The first four state slots remain queryWeight, keyWeight, valueWeight, outputWeight. Enabling input bias appends queryBias, keyBias, valueBias; enabling output bias appends outputBias; enabling dropout appends its non-trainable scalar probability. Disabled options allocate no state slots. All biases start at zero and consume no initialization seeds.

            Instances For