transformer_lens.model_bridge.composition_scores module

Tensor-like container for composition score results with layer-index metadata.

class transformer_lens.model_bridge.composition_scores.CompositionScores(scores: Tensor, layer_indices: List[int], head_labels: List[str])

Bases: object

Composition scores that behave like a tensor but carry layer-index metadata.

Delegates indexing, .shape, arithmetic, and torch.* functions to the underlying scores tensor via __torch_function__. On hybrid models where n_attn_layers < n_layers, layer_indices maps tensor position i to the original layer number.

scores

Upper-triangular composition score tensor.

layer_indices

Original layer numbers, e.g. [0, 2, 5].

head_labels

Labels matching scores dims, e.g. [“L0H0”, “L0H1”, …].

__getitem__(key)

property device: device

property dtype: dtype

property shape: Size