Struct TokenBatch 

pub struct TokenBatch {
    pub ids: Tensor,
    pub mask: Tensor,
    pub lengths: Vec<usize>,
}

A right-padded batch of token ids and the padding mask that goes with it.

Built once by pad_token_batch, then converted into whatever shape a given model’s forward pass wants. The four ported architectures disagree about that — ESM-2 takes a 1 = real token mask, AMPLIFY takes an additive pre-softmax bias, ESM-C and ESM-3 take a where_cond predicate — and deriving all three from one mask is what keeps that disagreement from becoming three subtly different notions of “padded”.

Fields

ids: Tensor

(batch, padded_len) token ids, dtype U32.

mask: Tensor

(batch, padded_len), dtype U8: 1 at a real token, 0 at padding.

lengths: Vec<usize>

Unpadded token count of each row, in input order.

Implementations

impl TokenBatch

pub fn padded_len(&self) -> usize

The length every row was padded to — the longest row.

pub fn mask_as(&self, dtype: DType) -> Result<Tensor>

The padding mask cast to dtype, for models that want it as a float.

Trait Implementations

impl Clone for TokenBatch

fn clone(&self) -> TokenBatch

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TokenBatch

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

Formats the value using the given formatter.