Struct AttentionVisualizer

pub struct AttentionVisualizer<F: Float + Debug> {
    pub num_heads: usize,
    pub sequence_length: usize,
    pub aggregation: AttentionAggregation,
    pub attention_cache: HashMap<String, ArrayD<F>>,
    pub target_layers: Vec<String>,
}

Attention visualizer for transformer models

Fields

num_heads: usize

Number of attention heads

sequence_length: usize

Sequence length

aggregation: AttentionAggregation

Aggregation strategy

attention_cache: HashMap<String, ArrayD<F>>

Cached attention weights

target_layers: Vec<String>

Layer names to visualize

Implementations

impl<F> AttentionVisualizer<F>

where F: Float + Debug + 'static + ScalarOperand + FromPrimitive + Sum + Clone + Copy,

pub fn new( num_heads: usize, sequence_length: usize, aggregation: AttentionAggregation, target_layers: Vec<String>, ) -> Self

Create a new attention visualizer

pub fn cache_attention_weights( &mut self, layer_name: String, attention_weights: ArrayD<F>, )

Cache attention weights for a layer

pub fn visualize_attention(&self, layer_name: &str) -> Result<ArrayD<F>>

Visualize attention patterns

pub fn aggregate_attention_heads( &self, attention_weights: &ArrayD<F>, ) -> Result<ArrayD<F>>

Aggregate attention across multiple heads

pub fn attention_rollout(&self) -> Result<ArrayD<F>>

Generate attention rollout visualization

pub fn visualize_attention_flow( &self, layer_name: &str, token_indices: &[usize], ) -> Result<Vec<f64>>

Visualize attention flow between tokens

Trait Implementations

impl<F: Clone + Float + Debug> Clone for AttentionVisualizer<F>

fn clone(&self) -> AttentionVisualizer<F>

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<F: Debug + Float + Debug> Debug for AttentionVisualizer<F>

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

Formats the value using the given formatter.