Coverage for transformer_lens/model_bridge/generalized_components/vision_embeddings.py: 100%

1"""Vision (ViT-style) embeddings bridge component. 

2 

3Wraps a HF `ViTEmbeddings` / `DeiTEmbeddings` module directly and forwards through 

4it unmodified. The patch-conv projection, CLS-token (and, for DeiT, distillation- 

5token) concatenation, and position-embedding addition are HF's math — we don't 

6reimplement any of it, we just point at the real module and hook its input/output. 

7 

8Why one class covers both ViT and DeiT: `ViTEmbeddings.forward` and 

9`DeiTEmbeddings.forward` have an *identical* signature 

10 (pixel_values, bool_masked_pos=None, interpolate_pos_encoding=False) 

11The only structural difference between them (DeiT prepends a distillation token 

12in addition to CLS, and sizes `position_embeddings` for +2 slots instead of +1) 

13lives entirely inside whichever module `set_original_component` resolves to — it's 

14invisible to this wrapper, so no DeiT-specific branching is needed here. 

15 

16Verified against the real generalized_components/base.py (not extrapolated): 

17`GeneralizedComponent.forward(*args, **kwargs)` already does exactly what we need 

18— it hooks the input via `self.hook_in`, casts it to match the wrapped module's 

19own parameter dtype (equivalent to HF's own `pixel_values.to(expected_dtype)` 

20"kept for BC" cast in ViTModel.forward — both land on the same dtype in the 

21overwhelming common case of a non-mixed-precision checkpoint), calls 

22`self.original_component(*args, **kwargs)`, and hooks the output via 

23`self.hook_out`. So this subclass only needs to guarantee `pixel_values` reaches 

24`super().forward()` *positionally* — the base class's own kwarg-name sniffing 

25list (`input`, `hidden_states`, `input_ids`, `query_input`, `x`, `inputs_embeds`) 

26doesn't include `pixel_values`, so if this component were ever called 

27all-keyword (`self.embeddings(pixel_values=x, ...)`, unlike the two HF call 

28sites we've confirmed) the base class's hook_in would silently never fire. 

29Re-emitting positionally here closes that gap regardless of how *we* were 

30called. 

31""" 

32 

33from typing import Any, Optional 

34 

35import torch 

36from torch import Tensor 

37 

38from transformer_lens.model_bridge.generalized_components.base import ( 

39 GeneralizedComponent, 

40) 

41 

42 

43class VisionEmbeddingsBridge(GeneralizedComponent): 

44 """Bridge for ViTEmbeddings / DeiTEmbeddings. 

45 

46 Point `name=` at the embeddings module directly, e.g.: 

47 VisionEmbeddingsBridge(name="embeddings") # bare ViTModel/DeiTModel 

48 VisionEmbeddingsBridge(name="vit.embeddings") # ViTForImageClassification 

49 VisionEmbeddingsBridge(name="deit.embeddings") # DeiTForImageClassification 

50 """ 

51 

52 def forward( 

53 self, 

54 pixel_values: Tensor, 

55 bool_masked_pos: Optional[torch.BoolTensor] = None, 

56 interpolate_pos_encoding: Optional[bool] = None, 

57 **kwargs: Any, 

58 ) -> Tensor: 

59 return super().forward( 

60 pixel_values, 

61 bool_masked_pos=bool_masked_pos, 

62 interpolate_pos_encoding=interpolate_pos_encoding, 

63 **kwargs, 

64 )