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

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 

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

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

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

21 

22_PAIRWISE_TEMP_BYTES = 64 * 1024 * 1024 

23 

24 

25@dataclass(frozen=True) 

26class SubspaceBasis: 

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

28 

29 Attributes: 

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

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

32 rank: Number of retained basis directions. 

33 measured_rank: Numerical rank before optional caller truncation. 

34 rtol: Effective relative rank tolerance. 

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

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

37 """ 

38 

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

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

41 rank: int 

42 measured_rank: int 

43 rtol: float 

44 threshold: float 

45 input_shape: Tuple[int, int] 

46 

47 @property 

48 def ambient_dim(self) -> int: 

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

50 return self.basis.shape[0] 

51 

52 

53@dataclass(frozen=True) 

54class ProjectionKernelResult: 

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

56 

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

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

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

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

61 rank_a: int 

62 rank_b: int 

63 ambient_dim: int 

64 

65 

66@dataclass(frozen=True) 

67class RandomSubspaceReference: 

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

69 

70 ambient_dim: int 

71 rank: int 

72 mean: float 

73 variance: float 

74 

75 

76@dataclass(frozen=True) 

77class AttentionHeadRef: 

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

79 

80 layer: int 

81 head: int 

82 role: AttentionRole 

83 kind: HeadKind 

84 

85 @property 

86 def label(self) -> str: 

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

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

89 

90 

91@dataclass(frozen=True) 

92class HeadAffinityPair: 

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

94 

95 source: AttentionHeadRef 

96 target: AttentionHeadRef 

97 score: float 

98 normalized: float 

99 

100 

101@dataclass(frozen=True) 

102class HeadAffinityResult: 

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

104 

105 Score tensors have shape 

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

107 tuples map tensor positions to original model block numbers. 

108 

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

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

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

112 """ 

113 

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

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

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

117 source_role: AttentionRole 

118 target_role: AttentionRole 

119 source_layer_indices: Tuple[int, ...] 

120 target_layer_indices: Tuple[int, ...] 

121 source_head_kind: HeadKind 

122 target_head_kind: HeadKind 

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

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

125 source_rank: int 

126 target_rank: int 

127 rank: Optional[int] 

128 rtol: float 

129 

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

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

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

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

134 

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

136 entries: List[HeadAffinityPair] = [] 

137 for source_layer, source_head, target_layer, target_head in ( 

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

139 ): 

140 source = AttentionHeadRef( 

141 layer=self.source_layer_indices[source_layer], 

142 head=source_head, 

143 role=self.source_role, 

144 kind=self.source_head_kind, 

145 ) 

146 target = AttentionHeadRef( 

147 layer=self.target_layer_indices[target_layer], 

148 head=target_head, 

149 role=self.target_role, 

150 kind=self.target_head_kind, 

151 ) 

152 entries.append( 

153 HeadAffinityPair( 

154 source=source, 

155 target=target, 

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

157 normalized=float( 

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

159 ), 

160 ) 

161 ) 

162 

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

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

165 return ( 

166 -value, 

167 pair.source.layer, 

168 pair.source.head, 

169 pair.target.layer, 

170 pair.target.head, 

171 ) 

172 

173 entries.sort(key=sort_key) 

174 return entries[:k] 

175 

176 

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

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

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

180 

181 

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

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

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

185 

186 

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

188 if rtol is None: 

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

190 storage_epsilon = torch.finfo(dtype).eps 

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

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

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

194 value = float(rtol) 

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

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

197 return value 

198 

199 

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

201 if rank is None: 

202 return None 

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

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

205 if not 1 <= rank <= max_rank: 

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

207 return rank 

208 

209 

210def orthonormal_subspace( 

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

212 *, 

213 rank: Optional[int] = None, 

214 rtol: Optional[float] = None, 

215) -> SubspaceBasis: 

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

217 

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

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

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

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

222 measured rank. 

223 

224 Args: 

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

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

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

228 

229 Returns: 

230 Basis, complete singular spectrum, and rank metadata. 

231 

232 Raises: 

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

234 """ 

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

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

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

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

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

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

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

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

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

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

245 

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

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

248 compute_dtype = _compute_dtype(matrix.dtype) 

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

250 work = matrix.to(dtype=compute_dtype) 

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

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

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

254 if measured_rank == 0: 

255 raise ValueError( 

256 "matrix numerical rank is zero " 

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

258 ) 

259 if requested_rank is not None and requested_rank > measured_rank: 

260 raise ValueError( 

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

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

263 ) 

264 

265 selected_rank = measured_rank if requested_rank is None else requested_rank 

266 return SubspaceBasis( 

267 basis=left[:, :selected_rank], 

268 singular_values=singular_values, 

269 rank=selected_rank, 

270 measured_rank=measured_rank, 

271 rtol=effective_rtol, 

272 threshold=threshold, 

273 input_shape=input_shape, 

274 ) 

275 

276 

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

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

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

280 basis = value.basis 

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

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

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

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

285 singular_values = value.singular_values 

286 if ( 

287 not isinstance(singular_values, torch.Tensor) 

288 or singular_values.ndim != 1 

289 or not torch.is_floating_point(singular_values) 

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

291 ): 

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

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

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

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

296 not isinstance(value.input_shape, tuple) 

297 or len(value.input_shape) != 2 

298 or any( 

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

300 for dimension in value.input_shape 

301 ) 

302 ): 

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

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

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

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

307 isinstance(value.rank, bool) 

308 or not isinstance(value.rank, int) 

309 or value.rank < 1 

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

311 ): 

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

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

314 isinstance(value.measured_rank, bool) 

315 or not isinstance(value.measured_rank, int) 

316 or value.measured_rank < value.rank 

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

318 ): 

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

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

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

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

323 isinstance(value.rtol, bool) 

324 or not isinstance(value.rtol, Real) 

325 or not math.isfinite(value.rtol) 

326 or value.rtol < 0 

327 ): 

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

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

330 isinstance(value.threshold, bool) 

331 or not isinstance(value.threshold, Real) 

332 or not math.isfinite(value.threshold) 

333 or value.threshold < 0 

334 ): 

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

336 

337 

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

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

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

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

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

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

344 raise ValueError( 

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

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

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

348 ) 

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

350 

351 

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

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

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

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

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

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

358 raise ValueError( 

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

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

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

362 ) 

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

364 

365 

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

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

368 if matrix.device.type == "mps": 368 ↛ 369line 368 didn't jump to line 369 because the condition on line 368 was never true

369 return torch.linalg.svdvals(matrix.cpu()).to(device=matrix.device) 

370 return torch.linalg.svdvals(matrix) 

371 

372 

373def projection_kernel( 

374 subspace_a: SubspaceBasis, 

375 subspace_b: SubspaceBasis, 

376 *, 

377 check_orthonormal: bool = True, 

378) -> ProjectionKernelResult: 

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

380 

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

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

383 matrices. Principal angles are returned in radians. 

384 """ 

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

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

387 _validate_subspace(subspace_a, "subspace_a") 

388 _validate_subspace(subspace_b, "subspace_b") 

389 if subspace_a.ambient_dim != subspace_b.ambient_dim: 

390 raise ValueError( 

391 "subspaces must have equal ambient dimensions, got " 

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

393 ) 

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

395 raise ValueError( 

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

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

398 ) 

399 

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

401 dtype = _compute_dtype(dtype) 

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

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

404 if check_orthonormal: 

405 eps = torch.finfo(dtype).eps 

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

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

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

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

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

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

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

413 

414 overlap = first.T @ second 

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

416 cosines = _clamp_principal_cosines(_singular_values(overlap)) 

417 angles = torch.acos(cosines) 

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

419 return ProjectionKernelResult( 

420 score=score, 

421 normalized=score / denominator, 

422 cosines=cosines, 

423 angles=angles, 

424 rank_a=subspace_a.rank, 

425 rank_b=subspace_b.rank, 

426 ambient_dim=subspace_a.ambient_dim, 

427 ) 

428 

429 

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

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

432 

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

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

435 """ 

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

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

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

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

440 

441 mean = rank**2 / ambient_dim 

442 variance = ( 

443 2 

444 * rank**2 

445 * (ambient_dim - rank) ** 2 

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

447 ) 

448 return RandomSubspaceReference( 

449 ambient_dim=ambient_dim, 

450 rank=rank, 

451 mean=float(mean), 

452 variance=float(variance), 

453 ) 

454 

455 

456def _pairwise_projection_kernel( 

457 source_bases: torch.Tensor, 

458 target_bases: torch.Tensor, 

459 *, 

460 max_temp_bytes: int = _PAIRWISE_TEMP_BYTES, 

461) -> torch.Tensor: 

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

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

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

465 if source_bases.shape[1] != target_bases.shape[1]: 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 share an ambient dimension") 

467 if source_bases.device != target_bases.device or source_bases.dtype != target_bases.dtype: 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 a device and dtype") 

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

470 isinstance(max_temp_bytes, bool) 

471 or not isinstance(max_temp_bytes, int) 

472 or max_temp_bytes < 1 

473 ): 

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

475 

476 source_count, _, source_rank = source_bases.shape 

477 target_count, _, target_rank = target_bases.shape 

478 scores = torch.empty( 

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

480 ) 

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

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

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

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

485 

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

487 source_stop = min(source_start + source_tile, source_count) 

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

489 target_stop = min(target_start + target_tile, target_count) 

490 overlap = torch.einsum( 

491 "adr,bds->abrs", 

492 source_bases[source_start:source_stop], 

493 target_bases[target_start:target_stop], 

494 ) 

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

496 dim=(-2, -1) 

497 ) 

498 return scores 

499 

500 

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

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

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

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

505 

506 

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

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

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

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

511 attribute = f"W_{role}" 

512 try: 

513 matrix = getattr(attn, attribute) 

514 except NotImplementedError as error: 

515 raise NotImplementedError( 

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

517 ) from error 

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

519 raise ValueError( 

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

521 ) from error 

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

523 raise ValueError( 

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

525 ) 

526 if matrix.ndim != 3: 

527 raise ValueError( 

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

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

530 ) 

531 if role == "O": 

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

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

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

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

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

537 return matrix.detach() 

538 

539 

540def _validate_role_shapes( 

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

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

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

544 if min(expected) < 1: 

545 raise ValueError( 

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

547 f"dimensions, got shape {expected}" 

548 ) 

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

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

551 raise ValueError( 

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

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

554 ) 

555 return expected 

556 

557 

558def _extract_bases( 

559 matrices: Sequence[torch.Tensor], 

560 layers: Sequence[int], 

561 role: AttentionRole, 

562 *, 

563 selected_rank: int, 

564 rtol: float, 

565 dtype: torch.dtype, 

566 device: torch.device, 

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

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

569 layer_ranks: List[List[int]] = [] 

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

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

572 head_ranks: List[int] = [] 

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

574 try: 

575 subspace = orthonormal_subspace( 

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

577 ) 

578 except ValueError as error: 

579 raise ValueError( 

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

581 ) from error 

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

583 head_ranks.append(subspace.measured_rank) 

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

585 layer_ranks.append(head_ranks) 

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

587 

588 

589def attention_head_subspace_affinity( 

590 model: Any, 

591 *, 

592 source_role: str = "O", 

593 target_role: str, 

594 layer_order: str = "forward", 

595 rank: Optional[int] = None, 

596 rtol: Optional[float] = None, 

597) -> HeadAffinityResult: 

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

599 

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

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

602 attention blocks and report their original block indices. 

603 

604 Args: 

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

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

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

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

609 keeps every pair. 

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

611 column rank. 

612 rtol: Optional relative numerical-rank tolerance. 

613 

614 Returns: 

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

616 """ 

617 if source_role != "O": 

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

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

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

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

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

623 validated_source_role = cast(AttentionRole, source_role) 

624 validated_target_role = cast(AttentionRole, target_role) 

625 validated_layer_order = cast(LayerOrder, layer_order) 

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

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

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

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

630 if not attention_blocks: 

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

632 

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

634 source_matrices = [ 

635 _read_role_matrix(model, block, layer, validated_source_role) 

636 for layer, block in attention_blocks 

637 ] 

638 target_matrices = [ 

639 _read_role_matrix(model, block, layer, validated_target_role) 

640 for layer, block in attention_blocks 

641 ] 

642 source_heads, source_ambient, source_width = _validate_role_shapes( 

643 source_matrices, layer_indices, validated_source_role 

644 ) 

645 target_heads, target_ambient, target_width = _validate_role_shapes( 

646 target_matrices, layer_indices, validated_target_role 

647 ) 

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

649 raise ValueError( 

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

651 f"{source_ambient} and {target_ambient}" 

652 ) 

653 

654 matrices = source_matrices + target_matrices 

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

656 dtype = matrices[0].dtype 

657 for matrix in matrices[1:]: 

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

659 dtype = _compute_dtype(dtype) 

660 effective_rtol = _validate_rtol( 

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

662 ) 

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

664 source_rank = source_width if requested_rank is None else requested_rank 

665 target_rank = target_width if requested_rank is None else requested_rank 

666 

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

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

669 result_device = ( 

670 torch.device(configured_device) 

671 if configured_device is not None 

672 else source_matrices[0].device 

673 ) 

674 source_bases, source_ranks = _extract_bases( 

675 source_matrices, 

676 layer_indices, 

677 validated_source_role, 

678 selected_rank=source_rank, 

679 rtol=effective_rtol, 

680 dtype=dtype, 

681 device=result_device, 

682 ) 

683 target_bases, target_ranks = _extract_bases( 

684 target_matrices, 

685 layer_indices, 

686 validated_target_role, 

687 selected_rank=target_rank, 

688 rtol=effective_rtol, 

689 dtype=dtype, 

690 device=result_device, 

691 ) 

692 

693 layer_count = len(layer_indices) 

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

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

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

697 layer_count, source_heads, layer_count, target_heads 

698 ) 

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

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

701 if validated_layer_order == "forward": 

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

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