TorchLean API

NN.Spec.Module.Seq2seq

Seq2Seq inference wrapper as an Spec.Module #

The Seq2Seq spec model defines encoder/decoder math and differentiable training helpers. This file provides a small inference-oriented Spec.Module wrapper so it can be composed/exported.

def Spec.Module.seq2seq {α : Type} [TorchLean.Storage α] [Context α] {srcVocabSize tgtVocabSize embedDim hiddenDim srcSeqLen tgtSeqLen : } (m : Seq2SeqSpec α srcVocabSize tgtVocabSize embedDim hiddenDim) (startToken : Fin tgtVocabSize) :
Module α [srcSeqLen, srcVocabSize] [tgtSeqLen, tgtVocabSize]

Seq2Seq inference module wrapper (one-hot input, greedy decoding).

Instances For