Struct LMBatch 

pub struct LMBatch {
    pub batch_size: usize,
    pub seq_len: usize,
    /* private fields */
}

A batch of tokenized sequences for language model training.

Memory Layout (ALB-100)

For causal LM, target[i] = input[i+1] — storing both wastes 50%. LMBatch deduplicates by storing a single tokens buffer:

• Shared layout (stride = seq_len + 1): Used when all sequences are the same length (the production pre-tokenized path). Input and target are derived as overlapping slices with a 1-token offset.

• Split layout (stride = 0): Used when sequences have different lengths (padding breaks the shift invariant). Stores [input_ids..., target_ids...] concatenated, matching the legacy layout.

Fields

batch_size: usize

Batch size

seq_len: usize

Sequence length (tokens per input/target per batch item)

Implementations

impl LMBatch

pub fn from_sequences(sequences: &[Vec<u32>], pad_id: u32, eos_id: u32) -> Self

Create a new LM batch from token sequences.

For causal LM, targets are inputs shifted by 1:

input:  [BOS, A, B, C, D]
target: [A, B, C, D, EOS]

When all sequences have the same length, uses shared layout (ALB-100): one buffer of batch_size * (seq_len + 1) tokens, saving ~50% memory.

When sequences differ in length, falls back to split layout (padding breaks the overlap invariant at sequence boundaries).

pub fn single(input_ids: Vec<u32>, target_ids: Vec<u32>) -> Self

Create a batch from a single sequence pair (for testing).

Uses split layout since caller provides separate input/target vecs.

pub fn get_input(&self, batch_idx: usize) -> Option<&[u32]>

Get input IDs for a specific batch item.

pub fn get_target(&self, batch_idx: usize) -> Option<&[u32]>

Get target IDs for a specific batch item.

pub fn num_tokens(&self) -> usize

Total number of tokens in batch (input side).

pub fn has_tokens(&self) -> bool

Returns true if the token buffer is non-empty.

Trait Implementations

impl Clone for LMBatch

fn clone(&self) -> LMBatch

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for LMBatch

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

Formats the value using the given formatter.