Coverage for transformer_lens/model_bridge/exceptions.py: 100%

1"""Exceptions for TransformerBridge. 

2 

3This module contains custom exceptions used by the TransformerBridge. 

4""" 

5from __future__ import annotations 

6 

7import torch 

8 

9 

10class StopAtLayerException(Exception): 

11 """Exception raised when execution should stop at a specific layer. 

12 

13 This exception is used to implement stop_at_layer functionality by having 

14 blocks check if they should stop and raise this exception, which is then 

15 caught in the main forward pass. 

16 """ 

17 

18 def __init__(self, layer_output: torch.Tensor): 

19 """Initialize with the output tensor from the layer where we stopped. 

20 

21 Args: 

22 layer_output: The output tensor from the last layer that executed 

23 """ 

24 self.layer_output = layer_output 

25 super().__init__()