pub fn ple_ngram_ids(dims: &PleIdDims, tokens: &[u32]) -> Vec<Vec<u64>>Expand description
Row ids for every head, one row per token in tokens.
tokens must already be context ++ new, where context is the
context_len preceding tokens (EOS-filled at the start of a sequence).
Returns [tokens.len()][ngram_heads]; callers slice off the last
new.len() rows, exactly as the reference’s
torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:] does.