Struct AttentionCache 

pub struct AttentionCache { /* private fields */ }

Stores per-layer attention weights from a forward pass.

Each tensor has shape [batch, heads, seq_q, seq_k] — the post-softmax attention pattern for one layer.

Example

use candle_mi::AttentionCache;
use candle_core::{Device, Tensor};

let mut cache = AttentionCache::with_capacity(32);
// shape [batch=1, heads=8, seq=10, seq=10]
cache.push(Tensor::zeros((1, 8, 10, 10), candle_core::DType::F32, &Device::Cpu).unwrap());

// Query what position 5 attends to in layer 0
let attn_row = cache.attention_from_position(0, 5).unwrap();

Implementations

impl AttentionCache

pub fn with_capacity(n_layers: usize) -> Self

Create an empty cache with capacity for n_layers layers.

pub fn push(&mut self, pattern: Tensor)

Add an attention pattern for the next layer.

Shapes

• pattern: [batch, heads, seq_q, seq_k]

pub const fn n_layers(&self) -> usize

Number of cached layers.

pub const fn is_empty(&self) -> bool

Whether the cache is empty.

pub fn get_layer(&self, layer: usize) -> Option<&Tensor>

Get the raw attention tensor for a specific layer.

Shapes

• returns: [batch, heads, seq_q, seq_k]

pub fn attention_from_position( &self, layer: usize, position: usize, ) -> Result<Vec<f32>>

Get attention weights FROM a specific query position, averaged across heads.

Returns a vector of length seq_k representing how much position attends to every key position, averaged over all attention heads (batch index 0).

Shapes

• returns: [seq_k] as Vec<f32>

Errors

Returns MIError::Hook if the layer is not in the cache or the position is out of range.

pub fn attention_to_position( &self, layer: usize, position: usize, ) -> Result<Vec<f32>>

Get attention weights TO a specific key position, averaged across heads.

Returns a vector of length seq_q representing how much every query position attends to position, averaged over all attention heads (batch index 0).

Shapes

• returns: [seq_q] as Vec<f32>

Errors

Returns MIError::Hook if the layer is not in the cache or the position is out of range.

pub fn top_attended_positions( &self, layer: usize, from_position: usize, k: usize, ) -> Result<Vec<(usize, f32)>>

Get the top-k key positions that a given query position attends to most.

Returns up to k pairs of (key_position, weight) sorted by descending attention weight, averaged across heads (batch index 0).

Errors

Returns MIError::Hook if the layer is not in the cache or the position is out of range.

pub fn patterns(&self) -> &[Tensor]

All cached patterns as a slice.

Trait Implementations

impl Debug for AttentionCache

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

Formats the value using the given formatter.