Implementation:LaurentMazare Tch rs Precompute Freqs Cis

Knowledge Sources	tch-rs
Domains	NLP, Positional_Encoding
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete tool for precomputing rotary position embedding frequency tensors provided by the tch-rs LLaMA example.

Description

precompute_freqs_cis generates a tensor of shape [1, 1, CONTEXT_SIZE, n_elem/2, 2] containing interleaved cosine and sine values for all positions and frequency pairs. The frequencies follow the geometric sequence 1/10000^(2i/n_elem) and are precomputed via an outer product of position indices and frequency values. The result is moved to the model's device and reused across all attention layers.

Usage

Call once during model initialization. Pass the resulting tensor to each Llama::forward call.

Code Reference

Source Location

• Repository: tch-rs
• File: examples/llama/main.rs
• Lines: 273-286

Signature

fn precompute_freqs_cis(config: &Config) -> Tensor

Import

// Internal to examples/llama/main.rs
use tch::Tensor;

I/O Contract

Inputs

Name	Type	Required	Description
config	&Config	Yes	Model config providing n_embd and n_head

Outputs

Name	Type	Description
Tensor	Tensor	Shape [1, 1, CONTEXT_SIZE, n_elem/2, 2] with interleaved cos/sin values

Usage Examples

let config = Config::config_7b();
let freqs_cis = precompute_freqs_cis(&config).to_device(device);

// Pass to every forward call
let logits = llama.forward(&tokens, &freqs_cis);

Related Pages

Implements Principle

• Principle:LaurentMazare_Tch_rs_Rotary_Position_Embedding