Coverage for transformer_lens/tools/analysis/projection_kernel.py: 88%

1"""Projection Kernel utilities for comparing linear subspaces. 

2 

3The Projection Kernel (PK) between subspaces with orthonormal bases ``U`` and 

4``V`` is ``||U.T @ V||_F^2``. It is invariant to basis choices within either 

5subspace and equals the sum of squared principal-angle cosines. 

6""" 

7 

8from __future__ import annotations 

9 

10import math 

11from dataclasses import dataclass 

12from numbers import Real 

13from typing import Any, List, Literal, Optional, Sequence, Tuple, cast 

14 

15import torch 

16from jaxtyping import Bool, Float, Int 

17 

18from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( 

19 _linalg_on_cpu_if_mps, 

20) 

21 

22AttentionRole = Literal["Q", "K", "V", "O"] 

23LayerOrder = Literal["forward", "all"] 

24HeadKind = Literal["query", "kv"] 

25 

26_PAIRWISE_TEMP_BYTES = 64 * 1024 * 1024 

27 

28 

29@dataclass(frozen=True) 

30class SubspaceBasis: 

31 """An explicitly ranked orthonormal basis extracted from a matrix. 

32 

33 Attributes: 

34 basis: Orthonormal column-space basis, ``[ambient_dim, rank]``. 

35 singular_values: All reduced-SVD singular values, in descending order. 

36 rank: Number of retained basis directions. 

37 measured_rank: Numerical rank before optional caller truncation. 

38 rtol: Effective relative rank tolerance. 

39 threshold: Absolute singular-value threshold used for rank measurement. 

40 input_shape: Shape of the matrix from which the basis was extracted. 

41 """ 

42 

43 basis: Float[torch.Tensor, "ambient rank"] 

44 singular_values: Float[torch.Tensor, "spectrum"] 

45 rank: int 

46 measured_rank: int 

47 rtol: float 

48 threshold: float 

49 input_shape: Tuple[int, int] 

50 

51 @property 

52 def ambient_dim(self) -> int: 

53 """Dimension of the space containing the subspace.""" 

54 return self.basis.shape[0] 

55 

56 

57@dataclass(frozen=True) 

58class ProjectionKernelResult: 

59 """Projection Kernel score and its principal-angle decomposition.""" 

60 

61 score: Float[torch.Tensor, ""] 

62 normalized: Float[torch.Tensor, ""] 

63 cosines: Float[torch.Tensor, "principal_angle"] 

64 angles: Float[torch.Tensor, "principal_angle"] 

65 rank_a: int 

66 rank_b: int 

67 ambient_dim: int 

68 

69 

70@dataclass(frozen=True) 

71class RandomSubspaceReference: 

72 """Analytic PK moments for independent random equal-rank subspaces.""" 

73 

74 ambient_dim: int 

75 rank: int 

76 mean: float 

77 variance: float 

78 

79 

80@dataclass(frozen=True) 

81class AttentionHeadRef: 

82 """Structured identity for one attention-head weight subspace.""" 

83 

84 layer: int 

85 head: int 

86 role: AttentionRole 

87 kind: HeadKind 

88 

89 @property 

90 def label(self) -> str: 

91 """Return the conventional TransformerLens layer/head label.""" 

92 return f"L{self.layer}H{self.head}" 

93 

94 

95@dataclass(frozen=True) 

96class HeadAffinityPair: 

97 """One ranked source-target head pair.""" 

98 

99 source: AttentionHeadRef 

100 target: AttentionHeadRef 

101 score: float 

102 normalized: float 

103 

104 

105@dataclass(frozen=True) 

106class HeadAffinityResult: 

107 """Projection Kernel affinities between two attention-head roles. 

108 

109 Score tensors have shape 

110 ``[source_layer, source_head, target_layer, target_head]``. Layer index 

111 tuples map tensor positions to original model block numbers. 

112 

113 ``source_ranks`` and ``target_ranks`` are measured numerical ranks for each 

114 head before optional truncation. Scalar ``source_rank`` and ``target_rank`` 

115 are the retained basis widths used for their respective roles. 

116 """ 

117 

118 scores: Float[torch.Tensor, "source_layer source_head target_layer target_head"] 

119 normalized: Float[torch.Tensor, "source_layer source_head target_layer target_head"] 

120 valid_mask: Bool[torch.Tensor, "source_layer source_head target_layer target_head"] 

121 source_role: AttentionRole 

122 target_role: AttentionRole 

123 source_layer_indices: Tuple[int, ...] 

124 target_layer_indices: Tuple[int, ...] 

125 source_head_kind: HeadKind 

126 target_head_kind: HeadKind 

127 source_ranks: Int[torch.Tensor, "source_layer source_head"] 

128 target_ranks: Int[torch.Tensor, "target_layer target_head"] 

129 source_rank: int 

130 target_rank: int 

131 rank: Optional[int] 

132 rtol: float 

133 

134 def top_pairs(self, k: int = 20, *, normalized: bool = False) -> List[HeadAffinityPair]: 

135 """Return the highest-scoring valid pairs with deterministic tie order.""" 

136 if isinstance(k, bool) or not isinstance(k, int) or k < 1: 

137 raise ValueError(f"k must be a positive integer, got {k!r}") 

138 

139 selected_scores = self.normalized if normalized else self.scores 

140 entries: List[HeadAffinityPair] = [] 

141 for source_layer, source_head, target_layer, target_head in ( 141 ↛ 167line 141 didn't jump to line 167 because the loop on line 141 didn't complete

142 torch.nonzero(self.valid_mask, as_tuple=False).cpu().tolist() 

143 ): 

144 source = AttentionHeadRef( 

145 layer=self.source_layer_indices[source_layer], 

146 head=source_head, 

147 role=self.source_role, 

148 kind=self.source_head_kind, 

149 ) 

150 target = AttentionHeadRef( 

151 layer=self.target_layer_indices[target_layer], 

152 head=target_head, 

153 role=self.target_role, 

154 kind=self.target_head_kind, 

155 ) 

156 entries.append( 

157 HeadAffinityPair( 

158 source=source, 

159 target=target, 

160 score=float(self.scores[source_layer, source_head, target_layer, target_head]), 

161 normalized=float( 

162 self.normalized[source_layer, source_head, target_layer, target_head] 

163 ), 

164 ) 

165 ) 

166 

167 def sort_key(pair: HeadAffinityPair) -> Tuple[float, int, int, int, int]: 

168 value = pair.normalized if normalized else pair.score 

169 return ( 

170 -value, 

171 pair.source.layer, 

172 pair.source.head, 

173 pair.target.layer, 

174 pair.target.head, 

175 ) 

176 

177 entries.sort(key=sort_key) 

178 return entries[:k] 

179 

180 

181def _compute_dtype(dtype: torch.dtype) -> torch.dtype: 

182 """Return a dtype supported by stable SVD on common PyTorch backends.""" 

183 return torch.float64 if dtype == torch.float64 else torch.float32 

184 

185 

186def _rank_tolerance_dtype(dtypes: Sequence[torch.dtype]) -> torch.dtype: 

187 """Return the least precise storage dtype for a shared rank tolerance.""" 

188 return max(dtypes, key=lambda dtype: torch.finfo(dtype).eps) 

189 

190 

191def _validate_rtol(rtol: Optional[float], shape: Tuple[int, int], dtype: torch.dtype) -> float: 

192 if rtol is None: 

193 compute_epsilon = torch.finfo(_compute_dtype(dtype)).eps 

194 storage_epsilon = torch.finfo(dtype).eps 

195 return max(max(shape) * compute_epsilon, storage_epsilon) 

196 if isinstance(rtol, bool) or not isinstance(rtol, Real): 196 ↛ 197line 196 didn't jump to line 197 because the condition on line 196 was never true

197 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") 

198 value = float(rtol) 

199 if not math.isfinite(value) or value < 0: 

200 raise ValueError(f"rtol must be a finite non-negative real number, got {rtol!r}") 

201 return value 

202 

203 

204def _validate_rank(rank: Optional[int], max_rank: int) -> Optional[int]: 

205 if rank is None: 

206 return None 

207 if isinstance(rank, bool) or not isinstance(rank, int): 

208 raise ValueError(f"rank must be an integer or None, got {rank!r}") 

209 if not 1 <= rank <= max_rank: 

210 raise ValueError(f"rank must be between 1 and min(matrix.shape)={max_rank}, got {rank}") 

211 return rank 

212 

213 

214def orthonormal_subspace( 

215 matrix: Float[torch.Tensor, "ambient width"], 

216 *, 

217 rank: Optional[int] = None, 

218 rtol: Optional[float] = None, 

219) -> SubspaceBasis: 

220 """Extract an explicitly ranked orthonormal column-space basis. 

221 

222 Low-precision inputs are promoted to float32 before the reduced SVD. With no 

223 explicit ``rtol``, numerical rank uses the larger of the compute-SVD error 

224 scale and one input-storage epsilon, relative to the largest singular value. 

225 An explicit ``rank`` truncates the measured subspace but may not exceed its 

226 measured rank. 

227 

228 Args: 

229 matrix: Finite floating-point matrix with shape ``[ambient_dim, width]``. 

230 rank: Optional number of leading singular directions to retain. 

231 rtol: Optional non-negative relative singular-value threshold. 

232 

233 Returns: 

234 Basis, complete singular spectrum, and rank metadata. 

235 

236 Raises: 

237 ValueError: If the matrix or rank policy is invalid. 

238 """ 

239 if not isinstance(matrix, torch.Tensor): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true

240 raise ValueError(f"matrix must be a torch.Tensor, got {type(matrix).__name__}") 

241 if matrix.ndim != 2: 241 ↛ 242line 241 didn't jump to line 242 because the condition on line 241 was never true

242 raise ValueError(f"matrix must be two-dimensional, got shape {tuple(matrix.shape)}") 

243 if matrix.shape[0] == 0 or matrix.shape[1] == 0: 

244 raise ValueError(f"matrix dimensions must be non-empty, got shape {tuple(matrix.shape)}") 

245 if not torch.is_floating_point(matrix): 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true

246 raise ValueError(f"matrix must have a real floating-point dtype, got {matrix.dtype}") 

247 if not bool(torch.isfinite(matrix).all()): 

248 raise ValueError("matrix must contain only finite values") 

249 

250 input_shape = (matrix.shape[0], matrix.shape[1]) 

251 requested_rank = _validate_rank(rank, min(input_shape)) 

252 compute_dtype = _compute_dtype(matrix.dtype) 

253 effective_rtol = _validate_rtol(rtol, input_shape, matrix.dtype) 

254 work = matrix.to(dtype=compute_dtype) 

255 left, singular_values, _ = torch.linalg.svd(work, full_matrices=False) 

256 threshold = float(singular_values[0].item()) * effective_rtol 

257 measured_rank = int((singular_values > threshold).sum().item()) 

258 if measured_rank == 0: 

259 raise ValueError( 

260 "matrix numerical rank is zero " 

261 f"for shape {input_shape}, dtype {matrix.dtype}, and threshold {threshold:.6g}" 

262 ) 

263 if requested_rank is not None and requested_rank > measured_rank: 

264 raise ValueError( 

265 f"requested rank {requested_rank} exceeds measured rank {measured_rank} " 

266 f"at threshold {threshold:.6g}" 

267 ) 

268 

269 selected_rank = measured_rank if requested_rank is None else requested_rank 

270 return SubspaceBasis( 

271 basis=left[:, :selected_rank], 

272 singular_values=singular_values, 

273 rank=selected_rank, 

274 measured_rank=measured_rank, 

275 rtol=effective_rtol, 

276 threshold=threshold, 

277 input_shape=input_shape, 

278 ) 

279 

280 

281def _validate_subspace(value: SubspaceBasis, name: str) -> None: 

282 if not isinstance(value, SubspaceBasis): 282 ↛ 283line 282 didn't jump to line 283 because the condition on line 282 was never true

283 raise ValueError(f"{name} must be a SubspaceBasis, got {type(value).__name__}") 

284 basis = value.basis 

285 if not isinstance(basis, torch.Tensor) or basis.ndim != 2: 285 ↛ 286line 285 didn't jump to line 286 because the condition on line 285 was never true

286 raise ValueError(f"{name}.basis must be a two-dimensional tensor") 

287 if not torch.is_floating_point(basis) or not bool(torch.isfinite(basis).all()): 

288 raise ValueError(f"{name}.basis must be finite and real floating-point") 

289 singular_values = value.singular_values 

290 if ( 

291 not isinstance(singular_values, torch.Tensor) 

292 or singular_values.ndim != 1 

293 or not torch.is_floating_point(singular_values) 

294 or not bool(torch.isfinite(singular_values).all()) 

295 ): 

296 raise ValueError(f"{name}.singular_values must be a finite floating-point vector") 

297 if singular_values.device != basis.device: 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true

298 raise ValueError(f"{name}.singular_values must be on the basis device") 

299 if ( 299 ↛ 307line 299 didn't jump to line 307 because the condition on line 299 was never true

300 not isinstance(value.input_shape, tuple) 

301 or len(value.input_shape) != 2 

302 or any( 

303 isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 1 

304 for dimension in value.input_shape 

305 ) 

306 ): 

307 raise ValueError(f"{name}.input_shape must contain two positive dimensions") 

308 if singular_values.shape[0] != min(value.input_shape): 308 ↛ 309line 308 didn't jump to line 309 because the condition on line 308 was never true

309 raise ValueError(f"{name}.singular_values length must equal min(input_shape)") 

310 if ( 310 ↛ 316line 310 didn't jump to line 316 because the condition on line 310 was never true

311 isinstance(value.rank, bool) 

312 or not isinstance(value.rank, int) 

313 or value.rank < 1 

314 or basis.shape[1] != value.rank 

315 ): 

316 raise ValueError(f"{name}.rank must equal its positive basis width") 

317 if ( 317 ↛ 323line 317 didn't jump to line 323 because the condition on line 317 was never true

318 isinstance(value.measured_rank, bool) 

319 or not isinstance(value.measured_rank, int) 

320 or value.measured_rank < value.rank 

321 or value.measured_rank > min(value.input_shape) 

322 ): 

323 raise ValueError(f"{name}.measured_rank must be at least its selected rank") 

324 if value.input_shape[0] != basis.shape[0]: 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true

325 raise ValueError(f"{name}.input_shape must share its basis ambient dimension") 

326 if ( 326 ↛ 332line 326 didn't jump to line 332 because the condition on line 326 was never true

327 isinstance(value.rtol, bool) 

328 or not isinstance(value.rtol, Real) 

329 or not math.isfinite(value.rtol) 

330 or value.rtol < 0 

331 ): 

332 raise ValueError(f"{name}.rtol must be finite and non-negative") 

333 if ( 333 ↛ 339line 333 didn't jump to line 339 because the condition on line 333 was never true

334 isinstance(value.threshold, bool) 

335 or not isinstance(value.threshold, Real) 

336 or not math.isfinite(value.threshold) 

337 or value.threshold < 0 

338 ): 

339 raise ValueError(f"{name}.threshold must be finite and non-negative") 

340 

341 

342def _clamp_projection_scores(scores: torch.Tensor, upper_bound: int) -> torch.Tensor: 

343 """Clamp roundoff-scale PK bound violations and reject larger violations.""" 

344 tolerance = 100 * torch.finfo(scores.dtype).eps * max(1, upper_bound) 

345 minimum = float(scores.detach().min().item()) 

346 maximum = float(scores.detach().max().item()) 

347 if minimum < -tolerance or maximum > upper_bound + tolerance: 

348 raise ValueError( 

349 "Projection Kernel score lies outside its theoretical bounds: " 

350 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, {upper_bound}] " 

351 f"within tolerance {tolerance:.6g}" 

352 ) 

353 return scores.clamp(min=0.0, max=float(upper_bound)) 

354 

355 

356def _clamp_principal_cosines(cosines: torch.Tensor) -> torch.Tensor: 

357 """Clamp roundoff-scale cosine violations and reject larger violations.""" 

358 tolerance = 100 * torch.finfo(cosines.dtype).eps * max(1, cosines.numel()) 

359 minimum = float(cosines.detach().min().item()) 

360 maximum = float(cosines.detach().max().item()) 

361 if minimum < -tolerance or maximum > 1 + tolerance: 

362 raise ValueError( 

363 "Principal-angle cosine lies outside its theoretical bounds: " 

364 f"observed [{minimum:.6g}, {maximum:.6g}], expected [0, 1] " 

365 f"within tolerance {tolerance:.6g}" 

366 ) 

367 return cosines.clamp(min=0.0, max=1.0) 

368 

369 

370def _singular_values(matrix: torch.Tensor) -> torch.Tensor: 

371 """Compute singular values with an explicit MPS CPU fallback.""" 

372 return _linalg_on_cpu_if_mps(torch.linalg.svdvals, matrix) 

373 

374 

375def projection_kernel( 

376 subspace_a: SubspaceBasis, 

377 subspace_b: SubspaceBasis, 

378 *, 

379 check_orthonormal: bool = True, 

380) -> ProjectionKernelResult: 

381 """Measure overlap between two explicitly extracted subspaces. 

382 

383 Raw PK lies in ``[0, min(rank_a, rank_b)]``. The normalized value is 

384 ``PK / sqrt(rank_a * rank_b)``, the cosine between the two projection 

385 matrices. Principal angles are returned in radians. 

386 """ 

387 if not isinstance(check_orthonormal, bool): 387 ↛ 388line 387 didn't jump to line 388 because the condition on line 387 was never true

388 raise ValueError(f"check_orthonormal must be a Boolean, got {check_orthonormal!r}") 

389 _validate_subspace(subspace_a, "subspace_a") 

390 _validate_subspace(subspace_b, "subspace_b") 

391 if subspace_a.ambient_dim != subspace_b.ambient_dim: 

392 raise ValueError( 

393 "subspaces must have equal ambient dimensions, got " 

394 f"{subspace_a.ambient_dim} and {subspace_b.ambient_dim}" 

395 ) 

396 if subspace_a.basis.device != subspace_b.basis.device: 396 ↛ 397line 396 didn't jump to line 397 because the condition on line 396 was never true

397 raise ValueError( 

398 "subspace bases must be on the same device, got " 

399 f"{subspace_a.basis.device} and {subspace_b.basis.device}" 

400 ) 

401 

402 dtype = torch.promote_types(subspace_a.basis.dtype, subspace_b.basis.dtype) 

403 dtype = _compute_dtype(dtype) 

404 first = subspace_a.basis.to(dtype=dtype) 

405 second = subspace_b.basis.to(dtype=dtype) 

406 if check_orthonormal: 

407 eps = torch.finfo(dtype).eps 

408 tolerance = 10 * max(first.shape[0], first.shape[1], second.shape[1]) * eps 

409 first_identity = torch.eye(first.shape[1], dtype=dtype, device=first.device) 

410 second_identity = torch.eye(second.shape[1], dtype=dtype, device=second.device) 

411 if not torch.allclose(first.T @ first, first_identity, rtol=tolerance, atol=tolerance): 

412 raise ValueError("subspace_a basis columns must be orthonormal") 

413 if not torch.allclose(second.T @ second, second_identity, rtol=tolerance, atol=tolerance): 413 ↛ 414line 413 didn't jump to line 414 because the condition on line 413 was never true

414 raise ValueError("subspace_b basis columns must be orthonormal") 

415 

416 overlap = first.T @ second 

417 score = _clamp_projection_scores(overlap.square().sum(), min(subspace_a.rank, subspace_b.rank)) 

418 cosines = _clamp_principal_cosines(_singular_values(overlap)) 

419 angles = torch.acos(cosines) 

420 denominator = math.sqrt(subspace_a.rank * subspace_b.rank) 

421 return ProjectionKernelResult( 

422 score=score, 

423 normalized=score / denominator, 

424 cosines=cosines, 

425 angles=angles, 

426 rank_a=subspace_a.rank, 

427 rank_b=subspace_b.rank, 

428 ambient_dim=subspace_a.ambient_dim, 

429 ) 

430 

431 

432def random_projection_kernel_moments(ambient_dim: int, rank: int) -> RandomSubspaceReference: 

433 """Return PK moments for independent Haar-distributed rank-``rank`` planes. 

434 

435 These idealized descriptive moments are not calibrated p-values for trained 

436 model weights, whose head subspaces are dependent and anisotropic. 

437 """ 

438 if isinstance(ambient_dim, bool) or not isinstance(ambient_dim, int) or ambient_dim < 2: 

439 raise ValueError(f"ambient_dim must be an integer at least 2, got {ambient_dim!r}") 

440 if isinstance(rank, bool) or not isinstance(rank, int) or not 1 <= rank <= ambient_dim: 

441 raise ValueError(f"rank must be an integer between 1 and {ambient_dim}, got {rank!r}") 

442 

443 mean = rank**2 / ambient_dim 

444 variance = ( 

445 2 

446 * rank**2 

447 * (ambient_dim - rank) ** 2 

448 / (ambient_dim**2 * (ambient_dim - 1) * (ambient_dim + 2)) 

449 ) 

450 return RandomSubspaceReference( 

451 ambient_dim=ambient_dim, 

452 rank=rank, 

453 mean=float(mean), 

454 variance=float(variance), 

455 ) 

456 

457 

458def _pairwise_projection_kernel( 

459 source_bases: torch.Tensor, 

460 target_bases: torch.Tensor, 

461 *, 

462 max_temp_bytes: int = _PAIRWISE_TEMP_BYTES, 

463) -> torch.Tensor: 

464 """Compute pairwise PK scores while bounding the overlap-tensor allocation.""" 

465 if source_bases.ndim != 3 or target_bases.ndim != 3: 465 ↛ 466line 465 didn't jump to line 466 because the condition on line 465 was never true

466 raise ValueError("source_bases and target_bases must be three-dimensional") 

467 if source_bases.shape[1] != target_bases.shape[1]: 467 ↛ 468line 467 didn't jump to line 468 because the condition on line 467 was never true

468 raise ValueError("source_bases and target_bases must share an ambient dimension") 

469 if source_bases.device != target_bases.device or source_bases.dtype != target_bases.dtype: 469 ↛ 470line 469 didn't jump to line 470 because the condition on line 469 was never true

470 raise ValueError("source_bases and target_bases must share a device and dtype") 

471 if ( 471 ↛ 476line 471 didn't jump to line 476 because the condition on line 471 was never true

472 isinstance(max_temp_bytes, bool) 

473 or not isinstance(max_temp_bytes, int) 

474 or max_temp_bytes < 1 

475 ): 

476 raise ValueError("max_temp_bytes must be a positive integer") 

477 

478 source_count, _, source_rank = source_bases.shape 

479 target_count, _, target_rank = target_bases.shape 

480 scores = torch.empty( 

481 source_count, target_count, dtype=source_bases.dtype, device=source_bases.device 

482 ) 

483 bytes_per_overlap = source_rank * target_rank * source_bases.element_size() 

484 pair_capacity = max(1, max_temp_bytes // bytes_per_overlap) 

485 source_tile = max(1, min(source_count, pair_capacity // max(1, target_count))) 

486 target_tile = max(1, min(target_count, pair_capacity // source_tile)) 

487 

488 for source_start in range(0, source_count, source_tile): 

489 source_stop = min(source_start + source_tile, source_count) 

490 for target_start in range(0, target_count, target_tile): 

491 target_stop = min(target_start + target_tile, target_count) 

492 overlap = torch.einsum( 

493 "adr,bds->abrs", 

494 source_bases[source_start:source_stop], 

495 target_bases[target_start:target_stop], 

496 ) 

497 scores[source_start:source_stop, target_start:target_stop] = overlap.square().sum( 

498 dim=(-2, -1) 

499 ) 

500 return scores 

501 

502 

503def _architecture_name(model: Any) -> str: 

504 cfg = getattr(model, "cfg", None) 

505 architecture = getattr(cfg, "original_architecture", None) 

506 return str(architecture) if architecture is not None else type(model).__name__ 

507 

508 

509def _read_role_matrix(model: Any, block: Any, layer: int, role: AttentionRole) -> torch.Tensor: 

510 attn = getattr(block, "attn", None) 

511 if attn is None: 511 ↛ 512line 511 didn't jump to line 512 because the condition on line 511 was never true

512 raise ValueError(f"attention block {layer} does not expose an attn component") 

513 attribute = f"W_{role}" 

514 try: 

515 matrix = getattr(attn, attribute) 

516 except NotImplementedError as error: 

517 raise NotImplementedError( 

518 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" 

519 ) from error 

520 except (AttributeError, RuntimeError, ValueError) as error: 

521 raise ValueError( 

522 f"{_architecture_name(model)} cannot expose role {role} at layer {layer}: {error}" 

523 ) from error 

524 if not isinstance(matrix, torch.Tensor): 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true

525 raise ValueError( 

526 f"role {role} at layer {layer} must be a tensor, got {type(matrix).__name__}" 

527 ) 

528 if matrix.ndim != 3: 

529 raise ValueError( 

530 f"role {role} at layer {layer} expected a three-dimensional per-head weight, " 

531 f"got shape {tuple(matrix.shape)}" 

532 ) 

533 if role == "O": 

534 matrix = matrix.transpose(-2, -1) 

535 if not torch.is_floating_point(matrix): 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true

536 raise ValueError(f"role {role} at layer {layer} must have a floating-point dtype") 

537 if not bool(torch.isfinite(matrix).all()): 

538 raise ValueError(f"role {role} at layer {layer} must contain only finite values") 

539 return matrix.detach() 

540 

541 

542def _validate_role_shapes( 

543 matrices: Sequence[torch.Tensor], layers: Sequence[int], role: AttentionRole 

544) -> Tuple[int, int, int]: 

545 expected = (matrices[0].shape[0], matrices[0].shape[1], matrices[0].shape[2]) 

546 if min(expected) < 1: 

547 raise ValueError( 

548 f"role {role} at layer {layers[0]} must have non-empty head, d_model, and width " 

549 f"dimensions, got shape {expected}" 

550 ) 

551 for matrix, layer in zip(matrices[1:], layers[1:]): 

552 if tuple(matrix.shape) != expected: 

553 raise ValueError( 

554 f"role {role} at layer {layer} has shape {tuple(matrix.shape)}, expected " 

555 f"the consistent [heads, d_model, width] shape {expected}" 

556 ) 

557 return expected 

558 

559 

560def _extract_bases( 

561 matrices: Sequence[torch.Tensor], 

562 layers: Sequence[int], 

563 role: AttentionRole, 

564 *, 

565 selected_rank: int, 

566 rtol: float, 

567 dtype: torch.dtype, 

568 device: torch.device, 

569) -> Tuple[torch.Tensor, torch.Tensor]: 

570 layer_bases: List[torch.Tensor] = [] 

571 layer_ranks: List[List[int]] = [] 

572 for matrix, layer in zip(matrices, layers): 

573 head_bases: List[torch.Tensor] = [] 

574 head_ranks: List[int] = [] 

575 for head in range(matrix.shape[0]): 

576 try: 

577 subspace = orthonormal_subspace( 

578 matrix[head].to(dtype=dtype), rank=selected_rank, rtol=rtol 

579 ) 

580 except ValueError as error: 

581 raise ValueError( 

582 f"Could not extract role {role} at layer {layer}, head {head}: {error}" 

583 ) from error 

584 head_bases.append(subspace.basis.to(device=device)) 

585 head_ranks.append(subspace.measured_rank) 

586 layer_bases.append(torch.stack(head_bases)) 

587 layer_ranks.append(head_ranks) 

588 return torch.stack(layer_bases), torch.tensor(layer_ranks, dtype=torch.long, device=device) 

589 

590 

591def attention_head_subspace_affinity( 

592 model: Any, 

593 *, 

594 source_role: str = "O", 

595 target_role: str, 

596 layer_order: str = "forward", 

597 rank: Optional[int] = None, 

598 rtol: Optional[float] = None, 

599) -> HeadAffinityResult: 

600 """Compute OQ, OK, or OV Projection Kernel affinities for a TransformerBridge. 

601 

602 K/V axes preserve native key-value heads on grouped-query attention models; 

603 they are never expanded to query-head count. Hybrid models include only 

604 attention blocks and report their original block indices. 

605 

606 Args: 

607 model: A TransformerBridge exposing readable per-head attention weights. 

608 source_role: Source role; v1 supports only ``"O"``. 

609 target_role: One of ``"Q"``, ``"K"``, or ``"V"``. 

610 layer_order: ``"forward"`` keeps strict earlier-to-later pairs; ``"all"`` 

611 keeps every pair. 

612 rank: Optional common truncation rank. By default every head must be full 

613 column rank. 

614 rtol: Optional relative numerical-rank tolerance. 

615 

616 Returns: 

617 Affinity tensors, validity mask, original layer indices, and rank metadata. 

618 """ 

619 if source_role != "O": 

620 raise ValueError(f"source_role must be 'O' in v1, got {source_role!r}") 

621 if target_role not in ("Q", "K", "V"): 

622 raise ValueError(f"target_role must be one of ['Q', 'K', 'V'], got {target_role!r}") 

623 if layer_order not in ("forward", "all"): 

624 raise ValueError(f"layer_order must be one of ['forward', 'all'], got {layer_order!r}") 

625 validated_source_role = cast(AttentionRole, source_role) 

626 validated_target_role = cast(AttentionRole, target_role) 

627 validated_layer_order = cast(LayerOrder, layer_order) 

628 blocks_with = getattr(model, "blocks_with", None) 

629 if not callable(blocks_with): 629 ↛ 630line 629 didn't jump to line 630 because the condition on line 629 was never true

630 raise ValueError("model must be a TransformerBridge exposing blocks_with('attn')") 

631 attention_blocks = list(blocks_with("attn")) 

632 if not attention_blocks: 

633 raise ValueError("No attention layers found — cannot compute head subspace affinity.") 

634 

635 layer_indices = [int(layer) for layer, _ in attention_blocks] 

636 source_matrices = [ 

637 _read_role_matrix(model, block, layer, validated_source_role) 

638 for layer, block in attention_blocks 

639 ] 

640 target_matrices = [ 

641 _read_role_matrix(model, block, layer, validated_target_role) 

642 for layer, block in attention_blocks 

643 ] 

644 source_heads, source_ambient, source_width = _validate_role_shapes( 

645 source_matrices, layer_indices, validated_source_role 

646 ) 

647 target_heads, target_ambient, target_width = _validate_role_shapes( 

648 target_matrices, layer_indices, validated_target_role 

649 ) 

650 if source_ambient != target_ambient: 650 ↛ 651line 650 didn't jump to line 651 because the condition on line 650 was never true

651 raise ValueError( 

652 f"roles {validated_source_role} and {validated_target_role} must share d_model, got " 

653 f"{source_ambient} and {target_ambient}" 

654 ) 

655 

656 matrices = source_matrices + target_matrices 

657 tolerance_dtype = _rank_tolerance_dtype([matrix.dtype for matrix in matrices]) 

658 dtype = matrices[0].dtype 

659 for matrix in matrices[1:]: 

660 dtype = torch.promote_types(dtype, matrix.dtype) 

661 dtype = _compute_dtype(dtype) 

662 effective_rtol = _validate_rtol( 

663 rtol, (source_ambient, max(source_width, target_width)), tolerance_dtype 

664 ) 

665 requested_rank = _validate_rank(rank, min(source_ambient, source_width, target_width)) 

666 source_rank = source_width if requested_rank is None else requested_rank 

667 target_rank = target_width if requested_rank is None else requested_rank 

668 

669 cfg = getattr(model, "cfg", None) 

670 configured_device = getattr(cfg, "device", None) 

671 result_device = ( 

672 torch.device(configured_device) 

673 if configured_device is not None 

674 else source_matrices[0].device 

675 ) 

676 source_bases, source_ranks = _extract_bases( 

677 source_matrices, 

678 layer_indices, 

679 validated_source_role, 

680 selected_rank=source_rank, 

681 rtol=effective_rtol, 

682 dtype=dtype, 

683 device=result_device, 

684 ) 

685 target_bases, target_ranks = _extract_bases( 

686 target_matrices, 

687 layer_indices, 

688 validated_target_role, 

689 selected_rank=target_rank, 

690 rtol=effective_rtol, 

691 dtype=dtype, 

692 device=result_device, 

693 ) 

694 

695 layer_count = len(layer_indices) 

696 flat_source = source_bases.reshape(layer_count * source_heads, source_ambient, source_rank) 

697 flat_target = target_bases.reshape(layer_count * target_heads, target_ambient, target_rank) 

698 scores = _pairwise_projection_kernel(flat_source, flat_target).reshape( 

699 layer_count, source_heads, layer_count, target_heads 

700 ) 

701 scores = _clamp_projection_scores(scores, min(source_rank, target_rank)) 

702 normalized = scores / math.sqrt(source_rank * target_rank) 

703 if validated_layer_order == "forward": 

704 layer_tensor = torch.tensor(layer_indices, device=result_device) 

705 layer_mask = layer_tensor[:, None] < layer_tensor[None, :]