Coverage for transformer_lens/HookedEncoder.py: 87%

1"""Hooked Encoder. 

2 

3Contains a BERT style model. This is separate from :class:`transformer_lens.HookedTransformer` 

4because it has a significantly different architecture to e.g. GPT style transformers. 

5""" 

6 

7from __future__ import annotations 

8 

9import logging 

10import os 

11import warnings 

12from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union, cast, overload 

13 

14import torch 

15from einops import repeat 

16from jaxtyping import Float, Int 

17from transformers.models.auto.tokenization_auto import AutoTokenizer 

18from typing_extensions import Literal 

19 

20import transformer_lens.loading_from_pretrained as loading 

21from transformer_lens.ActivationCache import ActivationCache 

22from transformer_lens.components import ( 

23 MLP, 

24 BertBlock, 

25 BertEmbed, 

26 BertMLMHead, 

27 BertNSPHead, 

28 BertPooler, 

29 Unembed, 

30) 

31from transformer_lens.components.mlps.gated_mlp import GatedMLP 

32from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

33from transformer_lens.FactoredMatrix import FactoredMatrix 

34from transformer_lens.hook_points import HookPoint 

35from transformer_lens.HookedRootModule import HookedRootModule 

36from transformer_lens.utilities import TypedModuleList, devices 

37 

38T = TypeVar("T", bound="HookedEncoder") 

39 

40 

41class HookedEncoder(HookedRootModule): 

42 """ 

43 This class implements a BERT-style encoder using the components in ./components.py, with HookPoints on every interesting activation. It inherits from HookedRootModule. 

44 

45 Limitations: 

46 - The model does not include dropouts, which may lead to inconsistent results from training or fine-tuning. 

47 

48 Like HookedTransformer, it can have a pretrained Transformer's weights loaded via `.from_pretrained`. There are a few features you might know from HookedTransformer which are not yet supported: 

49 - There is no preprocessing (e.g. LayerNorm folding) when loading a pretrained model 

50 """ 

51 

52 blocks: TypedModuleList[BertBlock] 

53 

54 def __init__( 

55 self, 

56 cfg: Union[HookedTransformerConfig, Dict], 

57 tokenizer: Optional[Any] = None, 

58 move_to_device: bool = True, 

59 **kwargs: Any, 

60 ): 

61 super().__init__() 

62 warnings.warn( 

63 "HookedEncoder is deprecated and will be removed in 4.0. Use " 

64 "TransformerBridge.boot_transformers(...) instead.", 

65 DeprecationWarning, 

66 stacklevel=2, 

67 ) 

68 if isinstance(cfg, Dict): 68 ↛ 69line 68 didn't jump to line 69 because the condition on line 68 was never true

69 cfg = HookedTransformerConfig(**cfg) 

70 elif isinstance(cfg, str): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 raise ValueError( 

72 "Please pass in a config dictionary or HookedTransformerConfig object. If you want to load a pretrained model, use HookedEncoder.from_pretrained() instead." 

73 ) 

74 self.cfg = cfg 

75 

76 assert self.cfg.n_devices == 1, "Multiple devices not supported for HookedEncoder" 

77 if tokenizer is not None: 

78 self.tokenizer = tokenizer 

79 elif self.cfg.tokenizer_name is not None: 

80 huggingface_token = os.environ.get("HF_TOKEN", "") 

81 self.tokenizer = AutoTokenizer.from_pretrained( 

82 self.cfg.tokenizer_name, 

83 token=huggingface_token if len(huggingface_token) > 0 else None, 

84 ) 

85 else: 

86 self.tokenizer = None 

87 

88 if self.cfg.d_vocab == -1: 

89 # If we have a tokenizer, vocab size can be inferred from it. 

90 assert self.tokenizer is not None, "Must provide a tokenizer if d_vocab is not provided" 

91 self.cfg.d_vocab = max(self.tokenizer.vocab.values()) + 1 

92 if self.cfg.d_vocab_out == -1: 

93 self.cfg.d_vocab_out = self.cfg.d_vocab 

94 

95 self.embed = BertEmbed(self.cfg) 

96 self.blocks = TypedModuleList([BertBlock(self.cfg) for _ in range(self.cfg.n_layers)]) 

97 self.mlm_head = BertMLMHead(self.cfg) 

98 self.unembed = Unembed(self.cfg) 

99 self.nsp_head = BertNSPHead(self.cfg) 

100 self.pooler = BertPooler(self.cfg) 

101 

102 self.hook_full_embed = HookPoint() 

103 

104 if move_to_device: 

105 if self.cfg.device is None: 105 ↛ 106line 105 didn't jump to line 106 because the condition on line 105 was never true

106 raise ValueError("Cannot move to device when device is None") 

107 self.to(self.cfg.device) 

108 

109 self.setup() 

110 

111 def to_tokens( 

112 self, 

113 input: Union[str, List[str]], 

114 move_to_device: bool = True, 

115 truncate: bool = True, 

116 ) -> Tuple[ 

117 Int[torch.Tensor, "batch pos"], 

118 Int[torch.Tensor, "batch pos"], 

119 Int[torch.Tensor, "batch pos"], 

120 ]: 

121 """Converts a string to a tensor of tokens. 

122 Taken mostly from the HookedTransformer implementation, but does not support default padding 

123 sides or prepend_bos. 

124 Args: 

125 input (Union[str, List[str]]): The input to tokenize. 

126 move_to_device (bool): Whether to move the output tensor of tokens to the device the model lives on. Defaults to True 

127 truncate (bool): If the output tokens are too long, whether to truncate the output 

128 tokens to the model's max context window. Does nothing for shorter inputs. Defaults to 

129 True. 

130 """ 

131 

132 assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer" 

133 

134 encodings = self.tokenizer( 

135 input, 

136 return_tensors="pt", 

137 padding=True, 

138 truncation=truncate, 

139 max_length=self.cfg.n_ctx if truncate else None, 

140 ) 

141 

142 tokens = encodings.input_ids 

143 token_type_ids = encodings.token_type_ids 

144 attention_mask = encodings.attention_mask 

145 

146 if move_to_device: 146 ↛ 151line 146 didn't jump to line 151 because the condition on line 146 was always true

147 tokens = tokens.to(self.cfg.device) 

148 token_type_ids = token_type_ids.to(self.cfg.device) 

149 attention_mask = attention_mask.to(self.cfg.device) 

150 

151 return tokens, token_type_ids, attention_mask 

152 

153 def encoder_output( 

154 self, 

155 tokens: Int[torch.Tensor, "batch pos"], 

156 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

157 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

158 ) -> Float[torch.Tensor, "batch pos d_vocab"]: 

159 """Processes input through the encoder layers and returns the resulting residual stream. 

160 

161 Args: 

162 input: Input tokens as integers with shape (batch, position) 

163 token_type_ids: Optional binary ids indicating segment membership. 

164 Shape (batch_size, sequence_length). For example, with input 

165 "[CLS] Sentence A [SEP] Sentence B [SEP]", token_type_ids would be 

166 [0, 0, ..., 0, 1, ..., 1, 1] where 0 marks tokens from sentence A 

167 and 1 marks tokens from sentence B. 

168 one_zero_attention_mask: Optional binary mask of shape (batch_size, sequence_length) 

169 where 1 indicates tokens to attend to and 0 indicates tokens to ignore. 

170 Used primarily for handling padding in batched inputs. 

171 

172 Returns: 

173 resid: Final residual stream tensor of shape (batch, position, d_model) 

174 

175 Raises: 

176 AssertionError: If using string input without a tokenizer 

177 """ 

178 

179 if tokens.device.type != self.cfg.device: 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true

180 tokens = tokens.to(self.cfg.device) 

181 if one_zero_attention_mask is not None: 

182 one_zero_attention_mask = one_zero_attention_mask.to(self.cfg.device) 

183 

184 resid = self.hook_full_embed(self.embed(tokens, token_type_ids)) 

185 

186 large_negative_number = -torch.inf 

187 mask = ( 

188 repeat(1 - one_zero_attention_mask, "batch pos -> batch 1 1 pos") 

189 if one_zero_attention_mask is not None 

190 else None 

191 ) 

192 additive_attention_mask = ( 

193 torch.where(mask == 1, large_negative_number, 0) if mask is not None else None 

194 ) 

195 

196 for block in self.blocks: 

197 resid = block(resid, additive_attention_mask) 

198 

199 return resid 

200 

201 @overload 

202 def forward( 

203 self, 

204 input: Union[ 

205 str, 

206 List[str], 

207 Int[torch.Tensor, "batch pos"], 

208 ], 

209 return_type: Union[Literal["logits"], Literal["predictions"]], 

210 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

211 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

212 ) -> Union[Float[torch.Tensor, "batch pos d_vocab"], str, List[str]]: 

213 ... 

214 

215 @overload 

216 def forward( 

217 self, 

218 input: Union[ 

219 str, 

220 List[str], 

221 Int[torch.Tensor, "batch pos"], 

222 ], 

223 return_type: Literal[None], 

224 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

225 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

226 ) -> Optional[Union[Float[torch.Tensor, "batch pos d_vocab"], str, List[str]]]: 

227 ... 

228 

229 def forward( 

230 self, 

231 input: Union[ 

232 str, 

233 List[str], 

234 Int[torch.Tensor, "batch pos"], 

235 ], 

236 return_type: Optional[Union[Literal["logits"], Literal["predictions"]]] = "logits", 

237 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

238 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

239 ) -> Optional[Union[Float[torch.Tensor, "batch pos d_vocab"], str, List[str]]]: 

240 """Forward pass through the HookedEncoder. Performs Masked Language Modelling on the given input. 

241 

242 Args: 

243 input: The input to process. Can be one of: 

244 - str: A single text string 

245 - List[str]: A list of text strings 

246 - torch.Tensor: Input tokens as integers with shape (batch, position) 

247 return_type: Optional[str]: The type of output to return. Can be one of: 

248 - None: Return nothing, don't calculate logits 

249 - 'logits': Return logits tensor 

250 - 'predictions': Return human-readable predictions 

251 token_type_ids: Optional[torch.Tensor]: Binary ids indicating whether a token belongs 

252 to sequence A or B. For example, for two sentences: 

253 "[CLS] Sentence A [SEP] Sentence B [SEP]", token_type_ids would be 

254 [0, 0, ..., 0, 1, ..., 1, 1]. `0` represents tokens from Sentence A, 

255 `1` from Sentence B. If not provided, BERT assumes a single sequence input. 

256 This parameter gets inferred from the tokenizer if input is a string or list of strings. 

257 Shape is (batch_size, sequence_length). 

258 one_zero_attention_mask: Optional[torch.Tensor]: A binary mask which indicates 

259 which tokens should be attended to (1) and which should be ignored (0). 

260 Primarily used for padding variable-length sentences in a batch. 

261 For instance, in a batch with sentences of differing lengths, shorter 

262 sentences are padded with 0s on the right. If not provided, the model 

263 assumes all tokens should be attended to. 

264 This parameter gets inferred from the tokenizer if input is a string or list of strings. 

265 Shape is (batch_size, sequence_length). 

266 

267 Returns: 

268 Optional[torch.Tensor]: Depending on return_type: 

269 - None: Returns None if return_type is None 

270 - torch.Tensor: Returns logits if return_type is 'logits' (or if return_type is not explicitly provided) 

271 - Shape is (batch_size, sequence_length, d_vocab) 

272 - str or List[str]: Returns predicted words for masked tokens if return_type is 'predictions'. 

273 Returns a list of strings if input is a list of strings, otherwise a single string. 

274 

275 Raises: 

276 AssertionError: If using string input without a tokenizer 

277 """ 

278 

279 if isinstance(input, str) or isinstance(input, list): 

280 assert self.tokenizer is not None, "Must provide a tokenizer if input is a string" 

281 tokens, token_type_ids_from_tokenizer, attention_mask = self.to_tokens(input) 

282 

283 # If token_type_ids or attention mask are not provided, use the ones from the tokenizer 

284 token_type_ids = ( 

285 token_type_ids_from_tokenizer if token_type_ids is None else token_type_ids 

286 ) 

287 one_zero_attention_mask = ( 

288 attention_mask if one_zero_attention_mask is None else one_zero_attention_mask 

289 ) 

290 

291 else: 

292 tokens = input 

293 

294 resid = self.encoder_output(tokens, token_type_ids, one_zero_attention_mask) 

295 

296 # MLM requires an unembedding step 

297 resid = self.mlm_head(resid) 

298 logits = self.unembed(resid) 

299 

300 if return_type == "predictions": 

301 assert ( 

302 self.tokenizer is not None 

303 ), "Must have a tokenizer to use return_type='predictions'" 

304 # Get predictions for masked tokens 

305 logprobs = logits[tokens == self.tokenizer.mask_token_id].log_softmax(dim=-1) 

306 predictions = self.tokenizer.decode(logprobs.argmax(dim=-1)) 

307 

308 # If input was a list of strings, split predictions into a list 

309 if " " in predictions: 309 ↛ 310line 309 didn't jump to line 310 because the condition on line 309 was never true

310 predictions = predictions.split(" ") 

311 predictions = [f"Prediction {i}: {p}" for i, p in enumerate(predictions)] 

312 return predictions 

313 

314 elif return_type == None: 314 ↛ 315line 314 didn't jump to line 315 because the condition on line 314 was never true

315 return None 

316 

317 return logits 

318 

319 @overload 

320 def run_with_cache( 

321 self, *model_args: Any, return_cache_object: Literal[True] = True, **kwargs: Any 

322 ) -> Tuple[Float[torch.Tensor, "batch pos d_vocab"], ActivationCache]: 

323 ... 

324 

325 @overload 

326 def run_with_cache( 

327 self, *model_args: Any, return_cache_object: Literal[False], **kwargs: Any 

328 ) -> Tuple[Float[torch.Tensor, "batch pos d_vocab"], Dict[str, torch.Tensor]]: 

329 ... 

330 

331 def run_with_cache( 

332 self, 

333 *model_args: Any, 

334 return_cache_object: bool = True, 

335 remove_batch_dim: bool = False, 

336 **kwargs: Any, 

337 ) -> Tuple[ 

338 Float[torch.Tensor, "batch pos d_vocab"], 

339 Union[ActivationCache, Dict[str, torch.Tensor]], 

340 ]: 

341 """ 

342 Wrapper around run_with_cache in HookedRootModule. If return_cache_object is True, this will return an ActivationCache object, with a bunch of useful HookedTransformer specific methods, otherwise it will return a dictionary of activations as in HookedRootModule. This function was copied directly from HookedTransformer. 

343 """ 

344 out, cache_dict = super().run_with_cache( 

345 *model_args, remove_batch_dim=remove_batch_dim, **kwargs 

346 ) 

347 if return_cache_object: 347 ↛ 351line 347 didn't jump to line 351 because the condition on line 347 was always true

348 cache = ActivationCache(cache_dict, self, has_batch_dim=not remove_batch_dim) 

349 return out, cache 

350 else: 

351 return out, cache_dict 

352 

353 def to( # type: ignore 

354 self, 

355 device_or_dtype: Union[torch.device, str, torch.dtype], 

356 print_details: bool = True, 

357 ): 

358 return devices.move_to_and_update_config(self, device_or_dtype, print_details) 

359 

360 def cuda(self: T, device: Optional[Union[int, torch.device]] = None) -> T: 

361 if isinstance(device, int): 

362 return self.to(f"cuda:{device}") 

363 elif device is None: 

364 return self.to("cuda") 

365 else: 

366 return self.to(device) 

367 

368 def cpu(self: T) -> T: 

369 return self.to("cpu") 

370 

371 def mps(self: T) -> T: 

372 """Warning: MPS may produce silently incorrect results. See #1178.""" 

373 return self.to(torch.device("mps")) 

374 

375 @classmethod 

376 def from_pretrained( 

377 cls, 

378 model_name: str, 

379 checkpoint_index: Optional[int] = None, 

380 checkpoint_value: Optional[int] = None, 

381 hf_model: Optional[Any] = None, 

382 device: Optional[str] = None, 

383 tokenizer: Optional[Any] = None, 

384 move_to_device: bool = True, 

385 dtype: torch.dtype = torch.float32, 

386 **from_pretrained_kwargs: Any, 

387 ) -> HookedEncoder: 

388 """Loads in the pretrained weights from huggingface. Currently supports loading weight from HuggingFace BertForMaskedLM. Unlike HookedTransformer, this does not yet do any preprocessing on the model.""" 

389 warnings.warn( 

390 "HookedEncoder.from_pretrained is deprecated and will be removed in 4.0. Use " 

391 "TransformerBridge.boot_transformers(...) instead.", 

392 DeprecationWarning, 

393 stacklevel=2, 

394 ) 

395 logging.warning( 

396 "Support for BERT in TransformerLens is currently experimental, until such a time when it has feature " 

397 "parity with HookedTransformer and has been tested on real research tasks. Until then, backward " 

398 "compatibility is not guaranteed. Please see the docs for information on the limitations of the current " 

399 "implementation." 

400 "\n" 

401 "If using BERT for interpretability research, keep in mind that BERT has some significant architectural " 

402 "differences to GPT. For example, LayerNorms are applied *after* the attention and MLP components, meaning " 

403 "that the last LayerNorm in a block cannot be folded." 

404 ) 

405 

406 assert not ( 

407 from_pretrained_kwargs.get("load_in_8bit", False) 

408 or from_pretrained_kwargs.get("load_in_4bit", False) 

409 ), "Quantization not supported" 

410 

411 if "torch_dtype" in from_pretrained_kwargs: 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true

412 dtype = from_pretrained_kwargs["torch_dtype"] 

413 

414 official_model_name = loading.get_official_model_name(model_name) 

415 

416 cfg = loading.get_pretrained_model_config( 

417 official_model_name, 

418 checkpoint_index=checkpoint_index, 

419 checkpoint_value=checkpoint_value, 

420 fold_ln=False, 

421 device=device, 

422 n_devices=1, 

423 dtype=dtype, 

424 **from_pretrained_kwargs, 

425 ) 

426 

427 state_dict = loading.get_pretrained_state_dict( 

428 official_model_name, cfg, hf_model, dtype=dtype, **from_pretrained_kwargs 

429 ) 

430 

431 model = cls(cfg, tokenizer, move_to_device=False) 

432 

433 model.load_state_dict(state_dict, strict=False) 

434 

435 if move_to_device: 435 ↛ 439line 435 didn't jump to line 439 because the condition on line 435 was always true

436 if cfg.device is not None: 436 ↛ 439line 436 didn't jump to line 439 because the condition on line 436 was always true

437 model.to(cfg.device) 

438 

439 print(f"Loaded pretrained model {model_name} into HookedEncoder") 

440 

441 return model 

442 

443 @property 

444 def W_U(self) -> Float[torch.Tensor, "d_model d_vocab"]: 

445 """ 

446 Convenience to get the unembedding matrix (ie the linear map from the final residual stream to the output logits) 

447 """ 

448 return self.unembed.W_U 

449 

450 @property 

451 def b_U(self) -> Float[torch.Tensor, "d_vocab"]: 

452 """ 

453 Convenience to get the unembedding bias 

454 """ 

455 return self.unembed.b_U 

456 

457 @property 

458 def W_E(self) -> Float[torch.Tensor, "d_vocab d_model"]: 

459 """ 

460 Convenience to get the embedding matrix 

461 """ 

462 return self.embed.embed.W_E 

463 

464 @property 

465 def W_pos(self) -> Float[torch.Tensor, "n_ctx d_model"]: 

466 """ 

467 Convenience function to get the positional embedding. Only works on models with absolute positional embeddings! 

468 """ 

469 return self.embed.pos_embed.W_pos 

470 

471 @property 

472 def W_E_pos(self) -> Float[torch.Tensor, "d_vocab+n_ctx d_model"]: 

473 """ 

474 Concatenated W_E and W_pos. Used as a full (overcomplete) basis of the input space, useful for full QK and full OV circuits. 

475 """ 

476 return torch.cat([self.W_E, self.W_pos], dim=0) 

477 

478 @property 

479 def W_K(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]: 

480 """Stacks the key weights across all layers""" 

481 return torch.stack([block.attn.W_K for block in self.blocks], dim=0) 

482 

483 @property 

484 def W_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]: 

485 """Stacks the query weights across all layers""" 

486 return torch.stack([block.attn.W_Q for block in self.blocks], dim=0) 

487 

488 @property 

489 def W_V(self) -> Float[torch.Tensor, "n_layers n_heads d_model d_head"]: 

490 """Stacks the value weights across all layers""" 

491 return torch.stack([block.attn.W_V for block in self.blocks], dim=0) 

492 

493 @property 

494 def W_O(self) -> Float[torch.Tensor, "n_layers n_heads d_head d_model"]: 

495 """Stacks the attn output weights across all layers""" 

496 return torch.stack([block.attn.W_O for block in self.blocks], dim=0) 

497 

498 @property 

499 def W_in(self) -> Float[torch.Tensor, "n_layers d_model d_mlp"]: 

500 """Stacks the MLP input weights across all layers""" 

501 return torch.stack( 

502 [cast(Union[MLP, GatedMLP], block.mlp).W_in for block in self.blocks], dim=0 

503 ) 

504 

505 @property 

506 def W_out(self) -> Float[torch.Tensor, "n_layers d_mlp d_model"]: 

507 """Stacks the MLP output weights across all layers""" 

508 return torch.stack( 

509 [cast(Union[MLP, GatedMLP], block.mlp).W_out for block in self.blocks], dim=0 

510 ) 

511 

512 @property 

513 def b_K(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]: 

514 """Stacks the key biases across all layers""" 

515 return torch.stack([block.attn.b_K for block in self.blocks], dim=0) 

516 

517 @property 

518 def b_Q(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]: 

519 """Stacks the query biases across all layers""" 

520 return torch.stack([block.attn.b_Q for block in self.blocks], dim=0) 

521 

522 @property 

523 def b_V(self) -> Float[torch.Tensor, "n_layers n_heads d_head"]: 

524 """Stacks the value biases across all layers""" 

525 return torch.stack([block.attn.b_V for block in self.blocks], dim=0) 

526 

527 @property 

528 def b_O(self) -> Float[torch.Tensor, "n_layers d_model"]: 

529 """Stacks the attn output biases across all layers""" 

530 return torch.stack([block.attn.b_O for block in self.blocks], dim=0) 

531 

532 @property 

533 def b_in(self) -> Float[torch.Tensor, "n_layers d_mlp"]: 

534 """Stacks the MLP input biases across all layers""" 

535 return torch.stack( 

536 [cast(Union[MLP, GatedMLP], block.mlp).b_in for block in self.blocks], dim=0 

537 ) 

538 

539 @property 

540 def b_out(self) -> Float[torch.Tensor, "n_layers d_model"]: 

541 """Stacks the MLP output biases across all layers""" 

542 return torch.stack( 

543 [cast(Union[MLP, GatedMLP], block.mlp).b_out for block in self.blocks], dim=0 

544 ) 

545 

546 @property 

547 def QK(self) -> FactoredMatrix: # [n_layers, n_heads, d_model, d_model] 

548 """Returns a FactoredMatrix object with the product of the Q and K matrices for each layer and head. 

549 Useful for visualizing attention patterns.""" 

550 return FactoredMatrix(self.W_Q, self.W_K.transpose(-2, -1)) 

551 

552 @property 

553 def OV(self) -> FactoredMatrix: # [n_layers, n_heads, d_model, d_model] 

554 """Returns a FactoredMatrix object with the product of the O and V matrices for each layer and head.""" 

555 return FactoredMatrix(self.W_V, self.W_O) 

556 

557 def all_head_labels(self) -> List[str]: 

558 """Returns a list of strings with the format "L{l}H{h}", where l is the layer index and h is the head index.""" 

559 return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)]