Struct std::JiT

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

struct JiT {
    &FloatTensor WPatchEmbed;
    struct Vec<struct JiTBlock> blocks;
}

A JiT model: the patch embedding and a stack of JiTBlocks. patches: [seqLen, patchDim]; cond: [1, dModel], broadcast-added to every token's embedding once, up front.

Fields

§&FloatTensor WPatchEmbedsource

[patchDim, dModel]

§struct Vec<struct JiTBlock> blockssource

Methods

§&FloatTensor forward(const &FloatTensor patches, const &FloatTensor cond) constsource · definition

-> [seqLen, dModel]

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

Parameterized (module.h): the embedding, then each block's weights.

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

§void drop()source · definition

§void free() consumesource · definition

= drop()

§

static struct JiT new(
    unsigned long patchDim,
    unsigned long dModel,
    unsigned long dHidden,
    unsigned long numHeads,
    unsigned long depth,
)

source · definition