Struct VocabNetwork 

pub struct VocabNetwork<M: Module> {
    pub embedding: Embedding,
    pub layers: Layers<M>,
    pub norm_f: RmsNorm,
    pub lm_head: Option<Linear>,
}

A complete autoregressive language model over a token vocabulary: Embedding (vocab → d_model) → Layers<M> → norm_f → LM head (d_model → vocab).

This is the token-LM counterpart of LatentNetwork; both are built on the shared Layers core. The only differences are the I/O boundary (a token Embedding and a vocab logit head, instead of two latent Linears) and a final pre-head RmsNorm.

The LM head is tied (lm_head = None, the transposed embedding weight is reused) or untied (a dedicated Linear); the vocabulary is rounded up to a multiple for GPU alignment (see VocabNetworkBuilder).

Fields

embedding: Embedding

Token embedding table, weight shape [padded_vocab, d_model].

layers: Layers<M>

The shared Mamba-x layer stack.

norm_f: RmsNorm

Final RMSNorm applied before the LM head (norm_f).

lm_head: Option<Linear>

Optional dedicated LM head. None ⇒ weight-tied (reuse embeddingᵀ).

Implementations

impl<M: MambaBlock> VocabNetwork<M>

where M::SsdPath: Clone,

pub fn forward( &self, x: Tensor<2, Int>, caches: Option<M::Caches>, ssd_path: M::SsdPath, ) -> (Tensor<3>, M::Caches)

Full-sequence pass: token IDs [batch, sequence] → logits [batch, sequence, padded_vocab].

pub fn step( &self, x: Tensor<1, Int>, caches: Option<M::Caches>, layers_own_index: Option<&mut usize>, layer_indices: Option<&mut Vec<usize>>, ) -> (Tensor<2>, M::Caches)

Single-token step: token IDs [batch] → logits [batch, padded_vocab].

The vocab network has no class tokens of its own (those would duplicate the layers’ class latents); it simply forwards the inner Layers cursors — layers_own_index (stack-level latents) and layer_indices (per-virtual-layer) — to Layers::step.

pub fn step_infinite(&self, x: Tensor<1, Int>) -> Tensor<2>

Stationary fixed point of the LM under a constant token: logits [batch, padded_vocab] after infinitely many repeats of x, no caches (see Layers::step_infinite).

pub fn step_n_approx( &self, x: Tensor<1, Int>, n: usize, caches: Option<M::Caches>, ) -> (Tensor<2>, M::Caches)

Approximate jump of n consecutive Self::step calls on the same constant token — see Layers::step_n_approx for the approximation contract.

fn apply_lm_head(&self, x: Tensor<3>) -> Tensor<3>

Project [batch, sequence, d_model] → [batch, sequence, padded_vocab] using the dedicated head, or the tied (transposed embedding) weight.

Trait Implementations

impl<M> AutodiffModule for VocabNetwork<M>

where M: AutodiffModule + ModuleDisplay + Module,

fn valid(&self) -> Self

Returns the same module, but on the inner backend without auto-differentiation.

fn from_inner(module: Self) -> Self

Wraps an inner module back into an auto-diff module.

impl<M> Clone for VocabNetwork<M>

where M: Module + ModuleDisplay + Module,

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<M: Debug + Module> Debug for VocabNetwork<M>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<M> Display for VocabNetwork<M>

where M: Module + ModuleDisplay + Module,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<M> Module for VocabNetwork<M>

where M: Module + ModuleDisplay + Module,

type Record = VocabNetworkRecord<M>

Type to save and load the module.

fn load_record(self, record: Self::Record) -> Self

Load the module state from a record.

fn into_record(self) -> Self::Record

Convert the module into a record containing the state.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules.

fn visit<Visitor: ModuleVisitor>(&self, visitor: &mut Visitor)

Visit each tensor parameter in the module with a visitor.

fn map<Mapper: ModuleMapper>(self, mapper: &mut Mapper) -> Self

Map each tensor parameter in the module with a mapper.

fn collect_devices(&self, devices: Devices) -> Devices

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn to_device(self, device: &Device) -> Self

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &Device) -> Self

Fork the module and all of its sub-modules to the given device.

fn devices(&self) -> Vec<Device>

Return all the devices found in the underneath module tree without duplicates.

fn no_grad(self) -> Self

Each tensor in the module tree will not require grad.

fn train(self) -> Self

where Self: AutodiffModule,

Move the module and all of its sub-modules to the autodiff backend.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

impl<M> ModuleDisplay for VocabNetwork<M>

where M: Module + ModuleDisplay + Module,

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<M> ModuleDisplayDefault for VocabNetwork<M>

where M: Module + ModuleDisplay + Module,

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.