Coverage for transformer_lens/utilities/slice.py: 94%

1"""Slice. 

2 

3This module contains the functionality for the Slice object 

4""" 

5 

6from __future__ import annotations 

7 

8import numbers 

9from typing import Any, List, Optional, Tuple, Union 

10 

11import numpy as np 

12import torch 

13 

14from .tensors import to_numpy 

15 

16# Type alias 

17SliceInput = Optional[ 

18 Union[ 

19 int, 

20 np.integer, 

21 Tuple[int,], 

22 Tuple[int, int], 

23 Tuple[int, int, int], 

24 List[int], 

25 torch.Tensor, 

26 np.ndarray, 

27 ] 

28] 

29"""An object that represents a slice input. It can be a tuple of integers or a slice object. 

30 

31An optional type alias for a slice input used in the `ActivationCache` module. 

32 

33A `SliceInput` can be one of the following types: 

34 - `int`: an integer representing a single position 

35 - `Tuple[int, int]`: a tuple of two integers representing a range of positions 

36 - `Tuple[int, int, int]`: a tuple of three integers representing a range of positions with a step size 

37 - `List[int]`: a list of integers representing multiple positions 

38 - `torch.Tensor`: a tensor containing a boolean mask or a list of indices to be selected from the input tensor. 

39 

40`SliceInput` is used in the `apply_ln_to_stack` method in the `ActivationCache` module. 

41""" 

42 

43 

44def to_python_int(value: Any) -> Optional[int]: 

45 """Return the Python int for a non-bool integer scalar (0-d tensor/ndarray too), else None.""" 

46 if isinstance(value, bool): 

47 return None 

48 if isinstance(value, numbers.Integral): 

49 return int(value) 

50 if isinstance(value, torch.Tensor): 

51 is_int_dtype = ( 

52 not value.dtype.is_floating_point 

53 and not value.dtype.is_complex 

54 and value.dtype != torch.bool 

55 ) 

56 if value.ndim == 0 and is_int_dtype: 

57 return int(value.item()) 

58 return None 

59 if isinstance(value, np.ndarray) and value.ndim == 0 and np.issubdtype(value.dtype, np.integer): 

60 return int(value.item()) 

61 return None 

62 

63 

64class Slice: 

65 """An object that represents a slice input. It can be a tuple of integers or a slice object. 

66 

67 We use a custom slice syntax because Python/Torch's don't let us reduce the number of dimensions: 

68 

69 Note that slicing with input_slice=None means do nothing, NOT add an extra dimension (use unsqueeze for that) 

70 

71 There are several modes: 

72 int - just index with that integer (decreases number of dimensions) 

73 slice - Input is a tuple converted to a slice ((k,) means :k, (k, m) means m:k, (k, m, n) means m:k:n) 

74 array - Input is a list or tensor or numpy array, converted to a numpy array, and we take the stack of values at those indices 

75 identity - Input is None, leave it unchanged. 

76 

77 Examples for dim=0: 

78 if input_slice=0, tensor -> tensor[0] 

79 elif input_slice = (1, 5), tensor -> tensor[1:5] 

80 elif input_slice = (1, 5, 2), tensor -> tensor[1:5:2] (ie indexing with [1, 3]) 

81 elif input_slice = [1, 4, 5], tensor -> tensor[[1, 4, 5]] (ie changing the first axis to have length 3, and taking the indices 1, 4, 5 out). 

82 elif input_slice is a Tensor, same as list - Tensor is assumed to be a 1D list of indices. 

83 """ 

84 

85 slice: Union[int, slice, np.ndarray] 

86 

87 def __init__( 

88 self, 

89 input_slice: SliceInput = None, 

90 ): 

91 """ 

92 Modular component for slicing tensors. Can be used to slice a tensor along a given dimension, or to index into a tensor along a given dimension. 

93 

94 Args: 

95 input_slice (SliceInput): The slice to apply. Can be an int, a tuple, a list, a torch.Tensor, or None. If None, do nothing. 

96 

97 Raises: 

98 ValueError: If the input_slice is not one of the above types. 

99 """ 

100 # np.integer and 0-d tensor scalars must index like ints, or downstream 

101 # mode=="int" checks miss them. 

102 int_input = to_python_int(input_slice) 

103 if int_input is not None: 

104 input_slice = int_input 

105 if isinstance(input_slice, tuple): 

106 self.slice = slice(*input_slice) 

107 self.mode = "slice" 

108 elif isinstance(input_slice, int): 

109 self.slice = input_slice 

110 self.mode = "int" 

111 elif isinstance(input_slice, slice): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true

112 self.slice = input_slice 

113 self.mode = "slice" 

114 elif type(input_slice) in [list, torch.Tensor, np.ndarray]: 

115 self.slice = to_numpy(input_slice) 

116 self.mode = "array" 

117 elif input_slice is None: 117 ↛ 121line 117 didn't jump to line 121 because the condition on line 117 was always true

118 self.slice = slice(None) 

119 self.mode = "identity" 

120 else: 

121 raise ValueError(f"Invalid input_slice {input_slice}") 

122 

123 def apply( 

124 self, 

125 tensor: torch.Tensor, 

126 dim: int = 0, 

127 ) -> torch.Tensor: 

128 """ 

129 Takes in a tensor and a slice, and applies the slice to the given dimension (supports positive and negative dimension syntax). Returns the sliced tensor. 

130 

131 Args: 

132 tensor (torch.Tensor): The tensor to slice. 

133 dim (int, optional): The dimension to slice along. Supports positive and negative dimension syntax. 

134 

135 Returns: 

136 torch.Tensor: The sliced tensor. 

137 """ 

138 ndim = tensor.ndim 

139 slices = [slice(None)] * ndim 

140 slices[dim] = self.slice # type: ignore 

141 return tensor[tuple(slices)] 

142 

143 def indices( 

144 self, 

145 max_ctx: Optional[int] = None, 

146 ) -> Union[np.ndarray, np.int32, np.int64]: 

147 """ 

148 Returns the indices when this slice is applied to an axis of size max_ctx. Returns them as a numpy array, for integer slicing it is eg array([4]) 

149 

150 Args: 

151 max_ctx (int, optional): The size of the axis to slice. Only used if the slice is not an integer. 

152 

153 Returns: 

154 np.ndarray: The indices that this slice will select. 

155 

156 Raises: 

157 ValueError: If the slice is not an integer and max_ctx is not specified. 

158 """ 

159 if self.mode == "int": 

160 return np.array([self.slice], dtype=np.int64) 

161 if max_ctx is None: 

162 raise ValueError("max_ctx must be specified if slice is not an integer") 

163 return np.arange(max_ctx, dtype=np.int64)[self.slice] 

164 

165 def __repr__( 

166 self, 

167 ) -> str: 

168 return f"Slice: {self.slice} Mode: {self.mode} " 

169 

170 @classmethod 

171 def unwrap( 

172 cls, 

173 slice_input: Union["Slice", SliceInput], 

174 ) -> "Slice": 

175 """ 

176 Takes a Slice-like input and converts it into a Slice, if it is not already. 

177 

178 Args: 

179 slice_input (Union[Slice, SliceInput]): The input to turn into a Slice. 

180 

181 Returns: 

182 Slice: A Slice object. 

183 """ 

184 if not isinstance(slice_input, Slice): 

185 int_input = to_python_int(slice_input) 

186 if int_input is not None: 

187 slice_input = int_input 

188 if isinstance( 

189 slice_input, int 

190 ): # slicing with an int collapses the dimension so this stops the pos dimension from collapsing 

191 slice_input = [slice_input] 

192 slice_input = Slice(slice_input) 

193 return slice_input