Skip to content

Tensor Activation Store¤

Tensor Activation Store.

TensorActivationStore ¤

Bases: ActivationStore

Tensor Activation Store.

Stores tensors in a (large) tensor of shape (item, neuron). Requires the number of activation vectors to be stored to be known in advance. Multiprocess safe.

Extends the torch.utils.data.Dataset class to provide a list-based activation store, with additional :meth:append and :meth:extend methods (the latter of which is non-blocking).

Examples: Create an empty activation dataset:

>>> import torch
>>> store = TensorActivationStore(max_items=1000, n_neurons=100, n_components=2)

Add a single activation vector to the dataset (for a component):

>>> store.append(torch.randn(100), component_idx=0)
>>> store.append(torch.randn(100), component_idx=1)
>>> len(store)
1

Add a [batch, neurons] activation tensor to the dataset:

>>> store.empty()
>>> batch = torch.randn(10, 100)
>>> store.extend(batch, component_idx=0)
>>> store.extend(batch, component_idx=1)
>>> len(store)
10

Shuffle the dataset before passing it to the DataLoader:

>>> store.shuffle() # Faster than using the DataLoader shuffle argument

Use the dataloader to iterate over the dataset:

>>> loader = torch.utils.data.DataLoader(store, shuffle=False, batch_size=2)
>>> next_item = next(iter(loader))
>>> next_item.shape
torch.Size([2, 2, 100])
Source code in sparse_autoencoder/activation_store/tensor_store.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
class TensorActivationStore(ActivationStore):
    """Tensor Activation Store.

    Stores tensors in a (large) tensor of shape (item, neuron). Requires the number of activation
    vectors to be stored to be known in advance. Multiprocess safe.

    Extends the `torch.utils.data.Dataset` class to provide a list-based activation store, with
    additional :meth:`append` and :meth:`extend` methods (the latter of which is non-blocking).

    Examples:
    Create an empty activation dataset:

        >>> import torch
        >>> store = TensorActivationStore(max_items=1000, n_neurons=100, n_components=2)

    Add a single activation vector to the dataset (for a component):

        >>> store.append(torch.randn(100), component_idx=0)
        >>> store.append(torch.randn(100), component_idx=1)
        >>> len(store)
        1

    Add a [batch, neurons] activation tensor to the dataset:

        >>> store.empty()
        >>> batch = torch.randn(10, 100)
        >>> store.extend(batch, component_idx=0)
        >>> store.extend(batch, component_idx=1)
        >>> len(store)
        10

    Shuffle the dataset **before passing it to the DataLoader**:

        >>> store.shuffle() # Faster than using the DataLoader shuffle argument

    Use the dataloader to iterate over the dataset:

        >>> loader = torch.utils.data.DataLoader(store, shuffle=False, batch_size=2)
        >>> next_item = next(iter(loader))
        >>> next_item.shape
        torch.Size([2, 2, 100])
    """

    _data: Float[Tensor, Axis.names(Axis.ITEMS, Axis.COMPONENT, Axis.INPUT_OUTPUT_FEATURE)]
    """Underlying Tensor Data Store."""

    _items_stored: list[int]
    """Number of items stored."""

    max_items: int
    """Maximum Number of Items to Store."""

    _n_components: int
    """Number of components"""

    @property
    def n_components(self) -> int:
        """Number of components."""
        return self._n_components

    @property
    def current_activations_stored_per_component(self) -> list[int]:
        """Number of activations stored per component."""
        return self._items_stored

    @validate_call(config={"arbitrary_types_allowed": True})
    def __init__(
        self,
        max_items: PositiveInt,
        n_neurons: PositiveInt,
        n_components: PositiveInt,
        device: torch.device | None = None,
    ) -> None:
        """Initialise the Tensor Activation Store.

        Args:
            max_items: Maximum number of items to store per component (individual activation
                vectors).
            n_neurons: Number of neurons in each activation vector.
            n_components: Number of components to store (i.e. number of source models).
            device: Device to store the activation vectors on.
        """
        self._n_components = n_components
        self._items_stored = [0] * n_components
        self._max_items = max_items
        self._data = torch.empty((max_items, n_components, n_neurons), device=device)

    def __len__(self) -> int:
        """Length Dunder Method.

        Returns the number of activation vectors per component in the dataset.

        Example:
            >>> import torch
            >>> store = TensorActivationStore(max_items=10_000_000, n_neurons=100, n_components=1)
            >>> store.append(torch.randn(100), component_idx=0)
            >>> store.append(torch.randn(100), component_idx=0)
            >>> len(store)
            2

        Returns:
            The number of activation vectors in the dataset.
        """
        # Min as this is the amount of activations that can be fetched by get_item
        return min(self.current_activations_stored_per_component)

    def __sizeof__(self) -> int:
        """Sizeof Dunder Method.

        Example:
            >>> import torch
            >>> store = TensorActivationStore(max_items=2, n_neurons=100, n_components=1)
            >>> store.__sizeof__() # Pre-allocated tensor of 2x100
            800

        Returns:
            The size of the underlying tensor in bytes.
        """
        return self._data.element_size() * self._data.nelement()

    def __getitem__(
        self, index: tuple[int, ...] | slice | int
    ) -> Float[Tensor, Axis.names(Axis.ANY)]:
        """Get Item Dunder Method.

        Examples:
            >>> import torch
            >>> store = TensorActivationStore(max_items=2, n_neurons=5, n_components=1)
            >>> store.append(torch.zeros(5), component_idx=0)
            >>> store.append(torch.ones(5), component_idx=0)
            >>> store[1, 0]
            tensor([1., 1., 1., 1., 1.])

        Args:
            index: The index of the tensor to fetch.

        Returns:
            The activation store item at the given index.
        """
        return self._data[index]

    def shuffle(self) -> None:
        """Shuffle the Data In-Place.

        This is much faster than using the shuffle argument on `torch.utils.data.DataLoader`.

        Example:
        >>> import torch
        >>> _seed = torch.manual_seed(42)
        >>> store = TensorActivationStore(max_items=10, n_neurons=1, n_components=1)
        >>> store.append(torch.tensor([0.]), component_idx=0)
        >>> store.append(torch.tensor([1.]), component_idx=0)
        >>> store.append(torch.tensor([2.]), component_idx=0)
        >>> store.shuffle()
        >>> [store[i, 0].item() for i in range(3)]
        [0.0, 2.0, 1.0]
        """
        # Generate a permutation of the indices for the active data
        perm = torch.randperm(len(self))

        # Use this permutation to shuffle the active data in-place
        self._data[: len(self)] = self._data[perm]

    def append(self, item: Float[Tensor, Axis.INPUT_OUTPUT_FEATURE], component_idx: int) -> None:
        """Add a single item to the store.

        Example:
        >>> import torch
        >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
        >>> store.append(torch.zeros(5), component_idx=0)
        >>> store.append(torch.ones(5), component_idx=0)
        >>> store[1, 0]
        tensor([1., 1., 1., 1., 1.])

        Args:
            item: The item to append to the dataset.
            component_idx: The component index to append the item to.

        Raises:
            IndexError: If there is no space remaining.
        """
        # Check we have space
        if self._items_stored[component_idx] + 1 > self._max_items:
            raise StoreFullError

        self._data[self._items_stored[component_idx], component_idx] = item.to(
            self._data.device,
        )
        self._items_stored[component_idx] += 1

    def extend(
        self,
        batch: Float[Tensor, Axis.names(Axis.BATCH, Axis.INPUT_OUTPUT_FEATURE)],
        component_idx: int,
    ) -> None:
        """Add a batch to the store.

        Examples:
        >>> import torch
        >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
        >>> store.extend(torch.zeros(2, 5), component_idx=0)
        >>> len(store)
        2

        Args:
            batch: The batch to append to the dataset.
            component_idx: The component index to append the batch to.

        Raises:
            IndexError: If there is no space remaining.
        """
        # Check we have space
        n_activation_tensors: int = batch.shape[0]
        if self._items_stored[component_idx] + n_activation_tensors > self._max_items:
            raise StoreFullError

        self._data[
            self._items_stored[component_idx] : self._items_stored[component_idx]
            + n_activation_tensors,
            component_idx,
        ] = batch.to(self._data.device)
        self._items_stored[component_idx] += n_activation_tensors

    def empty(self) -> None:
        """Empty the store.

        Example:
        >>> import torch
        >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
        >>> store.extend(torch.zeros(2, 5), component_idx=0)
        >>> len(store)
        2
        >>> store.empty()
        >>> len(store)
        0
        """
        # We don't need to zero the data, just reset the number of items stored
        self._items_stored = [0 for _ in self._items_stored]

current_activations_stored_per_component: list[int] property ¤

Number of activations stored per component.

max_items: int instance-attribute ¤

Maximum Number of Items to Store.

n_components: int property ¤

Number of components.

__getitem__(index) ¤

Get Item Dunder Method.

Examples:

>>> import torch
>>> store = TensorActivationStore(max_items=2, n_neurons=5, n_components=1)
>>> store.append(torch.zeros(5), component_idx=0)
>>> store.append(torch.ones(5), component_idx=0)
>>> store[1, 0]
tensor([1., 1., 1., 1., 1.])

Parameters:

Name Type Description Default
index tuple[int, ...] | slice | int

The index of the tensor to fetch.

 required

Returns:

Type Description
Float[Tensor, names(ANY)]

The activation store item at the given index.
Source code in sparse_autoencoder/activation_store/tensor_store.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def __getitem__(
    self, index: tuple[int, ...] | slice | int
) -> Float[Tensor, Axis.names(Axis.ANY)]:
    """Get Item Dunder Method.

    Examples:
        >>> import torch
        >>> store = TensorActivationStore(max_items=2, n_neurons=5, n_components=1)
        >>> store.append(torch.zeros(5), component_idx=0)
        >>> store.append(torch.ones(5), component_idx=0)
        >>> store[1, 0]
        tensor([1., 1., 1., 1., 1.])

    Args:
        index: The index of the tensor to fetch.

    Returns:
        The activation store item at the given index.
    """
    return self._data[index]

__init__(max_items, n_neurons, n_components, device=None) ¤

Initialise the Tensor Activation Store.

Parameters:

Name Type Description Default
max_items PositiveInt

Maximum number of items to store per component (individual activation vectors).

 required
n_neurons PositiveInt

Number of neurons in each activation vector.

 required
n_components PositiveInt

Number of components to store (i.e. number of source models).

 required
device device | None

Device to store the activation vectors on.

 None
Source code in sparse_autoencoder/activation_store/tensor_store.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
@validate_call(config={"arbitrary_types_allowed": True})
def __init__(
    self,
    max_items: PositiveInt,
    n_neurons: PositiveInt,
    n_components: PositiveInt,
    device: torch.device | None = None,
) -> None:
    """Initialise the Tensor Activation Store.

    Args:
        max_items: Maximum number of items to store per component (individual activation
            vectors).
        n_neurons: Number of neurons in each activation vector.
        n_components: Number of components to store (i.e. number of source models).
        device: Device to store the activation vectors on.
    """
    self._n_components = n_components
    self._items_stored = [0] * n_components
    self._max_items = max_items
    self._data = torch.empty((max_items, n_components, n_neurons), device=device)

__len__() ¤

Length Dunder Method.

Returns the number of activation vectors per component in the dataset.

Example

import torch store = TensorActivationStore(max_items=10_000_000, n_neurons=100, n_components=1) store.append(torch.randn(100), component_idx=0) store.append(torch.randn(100), component_idx=0) len(store) 2

Returns:

Type Description
int

The number of activation vectors in the dataset.
Source code in sparse_autoencoder/activation_store/tensor_store.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
def __len__(self) -> int:
    """Length Dunder Method.

    Returns the number of activation vectors per component in the dataset.

    Example:
        >>> import torch
        >>> store = TensorActivationStore(max_items=10_000_000, n_neurons=100, n_components=1)
        >>> store.append(torch.randn(100), component_idx=0)
        >>> store.append(torch.randn(100), component_idx=0)
        >>> len(store)
        2

    Returns:
        The number of activation vectors in the dataset.
    """
    # Min as this is the amount of activations that can be fetched by get_item
    return min(self.current_activations_stored_per_component)

__sizeof__() ¤

Sizeof Dunder Method.

Example

import torch store = TensorActivationStore(max_items=2, n_neurons=100, n_components=1) store.sizeof() # Pre-allocated tensor of 2x100 800

Returns:

Type Description
int

The size of the underlying tensor in bytes.
Source code in sparse_autoencoder/activation_store/tensor_store.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
def __sizeof__(self) -> int:
    """Sizeof Dunder Method.

    Example:
        >>> import torch
        >>> store = TensorActivationStore(max_items=2, n_neurons=100, n_components=1)
        >>> store.__sizeof__() # Pre-allocated tensor of 2x100
        800

    Returns:
        The size of the underlying tensor in bytes.
    """
    return self._data.element_size() * self._data.nelement()

append(item, component_idx) ¤

Add a single item to the store.

Example:

import torch store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1) store.append(torch.zeros(5), component_idx=0) store.append(torch.ones(5), component_idx=0) store[1, 0] tensor([1., 1., 1., 1., 1.])

Parameters:

Name Type Description Default
item Float[Tensor, INPUT_OUTPUT_FEATURE]

The item to append to the dataset.

 required
component_idx int

The component index to append the item to.

 required

Raises:

Type Description
IndexError

If there is no space remaining.
Source code in sparse_autoencoder/activation_store/tensor_store.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def append(self, item: Float[Tensor, Axis.INPUT_OUTPUT_FEATURE], component_idx: int) -> None:
    """Add a single item to the store.

    Example:
    >>> import torch
    >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
    >>> store.append(torch.zeros(5), component_idx=0)
    >>> store.append(torch.ones(5), component_idx=0)
    >>> store[1, 0]
    tensor([1., 1., 1., 1., 1.])

    Args:
        item: The item to append to the dataset.
        component_idx: The component index to append the item to.

    Raises:
        IndexError: If there is no space remaining.
    """
    # Check we have space
    if self._items_stored[component_idx] + 1 > self._max_items:
        raise StoreFullError

    self._data[self._items_stored[component_idx], component_idx] = item.to(
        self._data.device,
    )
    self._items_stored[component_idx] += 1

empty() ¤

Empty the store.

Example:

import torch store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1) store.extend(torch.zeros(2, 5), component_idx=0) len(store) 2 store.empty() len(store) 0

Source code in sparse_autoencoder/activation_store/tensor_store.py 
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
def empty(self) -> None:
    """Empty the store.

    Example:
    >>> import torch
    >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
    >>> store.extend(torch.zeros(2, 5), component_idx=0)
    >>> len(store)
    2
    >>> store.empty()
    >>> len(store)
    0
    """
    # We don't need to zero the data, just reset the number of items stored
    self._items_stored = [0 for _ in self._items_stored]

extend(batch, component_idx) ¤

Add a batch to the store.

Examples:

import torch store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1) store.extend(torch.zeros(2, 5), component_idx=0) len(store) 2

Parameters:

Name Type Description Default
batch Float[Tensor, names(BATCH, INPUT_OUTPUT_FEATURE)]

The batch to append to the dataset.

 required
component_idx int

The component index to append the batch to.

 required

Raises:

Type Description
IndexError

If there is no space remaining.
Source code in sparse_autoencoder/activation_store/tensor_store.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
def extend(
    self,
    batch: Float[Tensor, Axis.names(Axis.BATCH, Axis.INPUT_OUTPUT_FEATURE)],
    component_idx: int,
) -> None:
    """Add a batch to the store.

    Examples:
    >>> import torch
    >>> store = TensorActivationStore(max_items=10, n_neurons=5, n_components=1)
    >>> store.extend(torch.zeros(2, 5), component_idx=0)
    >>> len(store)
    2

    Args:
        batch: The batch to append to the dataset.
        component_idx: The component index to append the batch to.

    Raises:
        IndexError: If there is no space remaining.
    """
    # Check we have space
    n_activation_tensors: int = batch.shape[0]
    if self._items_stored[component_idx] + n_activation_tensors > self._max_items:
        raise StoreFullError

    self._data[
        self._items_stored[component_idx] : self._items_stored[component_idx]
        + n_activation_tensors,
        component_idx,
    ] = batch.to(self._data.device)
    self._items_stored[component_idx] += n_activation_tensors

shuffle() ¤

Shuffle the Data In-Place.

This is much faster than using the shuffle argument on torch.utils.data.DataLoader.

Example:

import torch _seed = torch.manual_seed(42) store = TensorActivationStore(max_items=10, n_neurons=1, n_components=1) store.append(torch.tensor([0.]), component_idx=0) store.append(torch.tensor([1.]), component_idx=0) store.append(torch.tensor([2.]), component_idx=0) store.shuffle() [store[i, 0].item() for i in range(3)] [0.0, 2.0, 1.0]

Source code in sparse_autoencoder/activation_store/tensor_store.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def shuffle(self) -> None:
    """Shuffle the Data In-Place.

    This is much faster than using the shuffle argument on `torch.utils.data.DataLoader`.

    Example:
    >>> import torch
    >>> _seed = torch.manual_seed(42)
    >>> store = TensorActivationStore(max_items=10, n_neurons=1, n_components=1)
    >>> store.append(torch.tensor([0.]), component_idx=0)
    >>> store.append(torch.tensor([1.]), component_idx=0)
    >>> store.append(torch.tensor([2.]), component_idx=0)
    >>> store.shuffle()
    >>> [store[i, 0].item() for i in range(3)]
    [0.0, 2.0, 1.0]
    """
    # Generate a permutation of the indices for the active data
    perm = torch.randperm(len(self))

    # Use this permutation to shuffle the active data in-place
    self._data[: len(self)] = self._data[perm]