Struct std::JiTBlock

#include <std/ml/transformer.h> source

struct JiTBlock {
    &FloatTensor Wq;
    &FloatTensor Wk;
    &FloatTensor Wv;
    &FloatTensor Wo;
    &FloatTensor W1;
    &FloatTensor b1;
    &FloatTensor W2;
    &FloatTensor b2;
    unsigned long numHeads;
}

Conditioning is added directly to the token embeddings once, up front (JiT::forward), rather than modulating every block's normalization the way DiT's adaLN does — the architectural distinction this implementation draws between the two. Each block is an ordinary pre-LN transformer sublayer pair:

h = LN(x); x = x + MHA(h@Wq,h@Wk,h@Wv,heads)@Wo
h2 = LN(x); x = x + relu(h2@W1+b1)@W2+b2

Fields

§&FloatTensor Wqsource

§&FloatTensor Wksource

§&FloatTensor Wvsource

§&FloatTensor Wosource

§&FloatTensor W1source

§&FloatTensor b1source

§&FloatTensor W2source

§&FloatTensor b2source

§unsigned long numHeadssource

Methods

§&FloatTensor forward(const &FloatTensor x) constsource · definition

§void collect_params(&Vec<&FloatTensor> out) constsource · definition

Parameterized (module.h): every weight tensor.

§void collect_state(&Vec<&FloatTensor> out) constsource · definition

§void drop()source · definition

§void free() consumesource · definition

= drop()

§

static struct JiTBlock new(
    unsigned long dModel,
    unsigned long dHidden,
    unsigned long numHeads,
)

source · definition