interp-engine 0.0.24__py3-none-any.whl

interp_engine-0.0.24.dist-info/METADATA

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: interp-engine
+Version: 0.0.24
+Summary: OpenAI and Neuronpedia's implementation of automated-interpretability, with some updates. Not officially affiliated with OpenAI.
+License-File: LICENSE
+Author: OpenAI, Neuronpedia
+Requires-Python: >=3.9,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: blobfile (>=2.1.1,<3.0.0)
+Requires-Dist: boostedblob (>=0.15.3,<0.16.0)
+Requires-Dist: httpx (>=0.27.0)
+Requires-Dist: numpy (>=1.24.0,<2.0.0)
+Requires-Dist: orjson (>=3.10.1,<4.0.0)
+Requires-Dist: scikit-learn (>=1.2.0,<2.0.0)
+Requires-Dist: tiktoken (>=0.6.0)

interp_engine-0.0.24.dist-info/RECORD

@@ -0,0 +1,28 @@
+neuron_explainer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+neuron_explainer/activations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+neuron_explainer/activations/activation_records.py,sha256=xRxbzO6zRQ24N6-ibySij9qbJo4CvyPwlj8drhD9KhQ,4751
+neuron_explainer/activations/activations.py,sha256=tLTx0xbdmh45YYq8JO36JIMuFDwynZjIuc6heNI8X7s,11282
+neuron_explainer/activations/attention_utils.py,sha256=Z7e125JoexP7TMkxw_kLW0_k8CIxKQ_rLNXaK0ZFbx4,5688
+neuron_explainer/activations/token_connections.py,sha256=mQSY1RCGG_RZr5WC3d1H7t_EZ73_oVV4HYIAM-0iKq4,2179
+neuron_explainer/api_client.py,sha256=fkTMR3ddvdKchtQPZmweL9aZmEHrxqbnvsMv7yQsKCk,7179
+neuron_explainer/azure.py,sha256=JljTuJTLFLCa80nxy7jFfGC2G2GOuHulqyUfXh9niSg,257
+neuron_explainer/explanations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+neuron_explainer/explanations/calibrated_simulator.py,sha256=ktW9ugS8vs0ZSlG1FsOj4ytLt2azFij8CSqgmqAMbcY,7967
+neuron_explainer/explanations/explainer.py,sha256=b03OqZPE-TFNwBToY25E6gMqxaY_Ycg--FNosc8B2Zg,122672
+neuron_explainer/explanations/explanations.py,sha256=XoBUX3kJoNONgkeBA--DtuF0erQGmzuyGDbyRuIq03Q,8867
+neuron_explainer/explanations/few_shot_examples.py,sha256=FY12WynQDtENi4ZwyaEIdBc26GwrtzSMx-wxSZryri8,88429
+neuron_explainer/explanations/prompt_builder.py,sha256=XS20hKKQVzbiyLckG8oERmCAPTW3Wu_uvjDPxdIugIw,4698
+neuron_explainer/explanations/puzzles.json,sha256=QjbopCB4hrYFM3aiYayQnnzLj6HYrq9plSHH9CJTrwc,52696
+neuron_explainer/explanations/puzzles.py,sha256=1Cnf_yPbCzz2EK3XRn5bryp7z6pTOSqk3yMc14A_-JY,2178
+neuron_explainer/explanations/scoring.py,sha256=bQPTQScthoy9VHOnztjfpLucCLpB_d260pRLdrODjTw,5888
+neuron_explainer/explanations/simulator.py,sha256=zvqs5A4N_kHfNdZBmVvSqvx531cMXm7sMqamZ2M70qY,45651
+neuron_explainer/explanations/test_explainer.py,sha256=kIPuIXHi0zWsvmHjNjGs-c19BXbmCfiqe94378Yh2O0,6603
+neuron_explainer/explanations/test_simulator.py,sha256=jlAAzUD_qufAxAIuaOJ54ZbcWs0jaQpGhIdpcnXQGp0,7605
+neuron_explainer/explanations/token_space_few_shot_examples.py,sha256=z5zR-TWF94yg1QD6WEBVfUQkZFo8WIBec8BDxDK19FI,4836
+neuron_explainer/fast_dataclasses/__init__.py,sha256=L_Swq2A_F2aKjEXY_srALjDoWI3pC1TwLJHIqBuUknc,147
+neuron_explainer/fast_dataclasses/fast_dataclasses.py,sha256=b9EGP7m2m1qZuC0wqzI-_kfus6jqxIW_f_fGVnFepCs,3043
+neuron_explainer/fast_dataclasses/test_fast_dataclasses.py,sha256=wxhS7S6tlOPNRKAUdX0sTSIhcMzty7T_5jVOip4m_DI,2712
+interp_engine-0.0.24.dist-info/METADATA,sha256=vfvJrjlGnEy4lDQ9e_xh_mrHtHdnUyjkIR-KtZqtgtc,904
+interp_engine-0.0.24.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+interp_engine-0.0.24.dist-info/licenses/LICENSE,sha256=Qplde_TQOQKYG75mHztmzHl3Zsnz7cZ8axqqM5XbkII,1065
+interp_engine-0.0.24.dist-info/RECORD,,

interp_engine-0.0.24.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.2.1
+Root-Is-Purelib: true
+Tag: py3-none-any

interp_engine-0.0.24.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2023 Superalignment, OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

neuron_explainer/activations/activation_records.py

@@ -0,0 +1,130 @@
+"""Utilities for formatting activation records into prompts."""
+
+import math
+from typing import Optional, Sequence
+
+from neuron_explainer.activations.activations import ActivationRecord
+
+UNKNOWN_ACTIVATION_STRING = "unknown"
+
+
+def relu(x: float) -> float:
+    return max(0.0, x)
+
+
+def calculate_max_activation(activation_records: Sequence[ActivationRecord]) -> float:
+    """Return the maximum activation value of the neuron across all the activation records."""
+    flattened = [
+        # Relu is used to assume any values less than 0 are indicating the neuron is in the resting
+        # state. This is a simplifying assumption that works with relu/gelu.
+        max(relu(x) for x in activation_record.activations)
+        for activation_record in activation_records
+    ]
+    return max(flattened)
+
+
+def normalize_activations(activation_record: list[float], max_activation: float) -> list[int]:
+    """Convert raw neuron activations to integers on the range [0, 10]."""
+    if max_activation <= 0:
+        return [0 for x in activation_record]
+    # Relu is used to assume any values less than 0 are indicating the neuron is in the resting
+    # state. This is a simplifying assumption that works with relu/gelu.
+    return [min(10, math.floor(10 * relu(x) / max_activation)) for x in activation_record]
+
+
+def _format_activation_record(
+    activation_record: ActivationRecord,
+    max_activation: float,
+    omit_zeros: bool,
+    hide_activations: bool = False,
+    start_index: int = 0,
+) -> str:
+    """Format neuron activations into a string, suitable for use in prompts."""
+    tokens = activation_record.tokens
+    normalized_activations = normalize_activations(activation_record.activations, max_activation)
+    if omit_zeros:
+        assert (not hide_activations) and start_index == 0, "Can't hide activations and omit zeros"
+        tokens = [
+            token for token, activation in zip(tokens, normalized_activations) if activation > 0
+        ]
+        normalized_activations = [x for x in normalized_activations if x > 0]
+
+    entries = []
+    assert len(tokens) == len(normalized_activations)
+    for index, token, activation in zip(range(len(tokens)), tokens, normalized_activations):
+        activation_string = str(int(activation))
+        if hide_activations or index < start_index:
+            activation_string = UNKNOWN_ACTIVATION_STRING
+        entries.append(f"{token}\t{activation_string}")
+    return "\n".join(entries)
+
+
+def format_activation_records(
+    activation_records: Sequence[ActivationRecord],
+    max_activation: float,
+    *,
+    omit_zeros: bool = False,
+    start_indices: Optional[list[int]] = None,
+    hide_activations: bool = False,
+) -> str:
+    """Format a list of activation records into a string."""
+    return (
+        "\n<start>\n"
+        + "\n<end>\n<start>\n".join(
+            [
+                _format_activation_record(
+                    activation_record,
+                    max_activation,
+                    omit_zeros=omit_zeros,
+                    hide_activations=hide_activations,
+                    start_index=0 if start_indices is None else start_indices[i],
+                )
+                for i, activation_record in enumerate(activation_records)
+            ]
+        )
+        + "\n<end>\n"
+    )
+
+
+def _format_tokens_for_simulation(tokens: Sequence[str]) -> str:
+    """
+    Format tokens into a string with each token marked as having an "unknown" activation, suitable
+    for use in prompts.
+    """
+    entries = []
+    for token in tokens:
+        entries.append(f"{token}\t{UNKNOWN_ACTIVATION_STRING}")
+    return "\n".join(entries)
+
+
+def format_sequences_for_simulation(
+    all_tokens: Sequence[Sequence[str]],
+) -> str:
+    """
+    Format a list of lists of tokens into a string with each token marked as having an "unknown"
+    activation, suitable for use in prompts.
+    """
+    return (
+        "\n<start>\n"
+        + "\n<end>\n<start>\n".join(
+            [_format_tokens_for_simulation(tokens) for tokens in all_tokens]
+        )
+        + "\n<end>\n"
+    )
+
+
+def non_zero_activation_proportion(
+    activation_records: Sequence[ActivationRecord], max_activation: float
+) -> float:
+    """Return the proportion of activation values that aren't zero."""
+    total_activations_count = sum(
+        [len(activation_record.activations) for activation_record in activation_records]
+    )
+    normalized_activations = [
+        normalize_activations(activation_record.activations, max_activation)
+        for activation_record in activation_records
+    ]
+    non_zero_activations_count = sum(
+        [len([x for x in activations if x != 0]) for activations in normalized_activations]
+    )
+    return non_zero_activations_count / total_activations_count

neuron_explainer/activations/activations.py

@@ -0,0 +1,311 @@
+# Dataclasses and enums for storing neuron-indexed information about activations. Also, related
+# helper functions.
+
+import math
+from dataclasses import dataclass, field
+from typing import List, Optional, Union
+
+import urllib.request
+import blobfile as bf
+import boostedblob as bbb
+from neuron_explainer.fast_dataclasses import FastDataclass, loads, register_dataclass
+from neuron_explainer.azure import standardize_azure_url
+
+
+@register_dataclass
+@dataclass
+class ActivationRecord(FastDataclass):
+    """Collated lists of tokens and their activations for a single neuron."""
+
+    tokens: List[str]
+    """Tokens in the text sequence, represented as strings."""
+    activations: List[float]
+    """Raw activation values for the neuron on each token in the text sequence."""
+    dfa_values: Optional[List[int]] = None
+    dfa_target_index: Optional[int] = None
+
+
+@register_dataclass
+@dataclass
+class NeuronId(FastDataclass):
+    """Identifier for a neuron in an artificial neural network."""
+
+    layer_index: int
+    """The index of layer the neuron is in. The first layer used during inference has index 0."""
+    neuron_index: int
+    """The neuron's index within in its layer. Indices start from 0 in each layer."""
+
+
+def _check_slices(
+    slices_by_split: dict[str, slice],
+    expected_num_values: int,
+) -> None:
+    """Assert that the slices are disjoint and fully cover the intended range."""
+    indices = set()
+    sum_of_slice_lengths = 0
+    n_splits = len(slices_by_split.keys())
+    for s in slices_by_split.values():
+        subrange = range(expected_num_values)[s]
+        sum_of_slice_lengths += len(subrange)
+        indices |= set(subrange)
+    assert (
+        sum_of_slice_lengths == expected_num_values
+    ), f"{sum_of_slice_lengths=} != {expected_num_values=}"
+    stride = n_splits
+    expected_indices = set.union(
+        *[
+            set(range(start_index, expected_num_values, stride))
+            for start_index in range(n_splits)
+        ]
+    )
+    assert indices == expected_indices, f"{indices=} != {expected_indices=}"
+
+
+def get_slices_for_splits(
+    splits: list[str],
+    num_activation_records_per_split: int,
+) -> dict[str, slice]:
+    """
+    Get equal-sized interleaved subsets for each of a list of splits, given the number of elements
+    to include in each split.
+    """
+
+    stride = len(splits)
+    num_activation_records_for_even_splits = num_activation_records_per_split * stride
+    slices_by_split = {
+        split: slice(split_index, num_activation_records_for_even_splits, stride)
+        for split_index, split in enumerate(splits)
+    }
+    _check_slices(
+        slices_by_split=slices_by_split,
+        expected_num_values=num_activation_records_for_even_splits,
+    )
+    return slices_by_split
+
+
+@dataclass
+class ActivationRecordSliceParams:
+    """How to select splits (train, valid, etc.) of activation records."""
+
+    n_examples_per_split: Optional[int]
+    """The number of examples to include in each split."""
+
+
+@register_dataclass
+@dataclass
+class NeuronRecord(FastDataclass):
+    """Neuron-indexed activation data, including summary stats and notable activation records."""
+
+    neuron_id: NeuronId
+    """Identifier for the neuron."""
+
+    random_sample: list[ActivationRecord] = field(default_factory=list)
+    """
+    Random activation records for this neuron. The random sample is independent from those used for
+    other neurons.
+    """
+    random_sample_by_quantile: Optional[list[list[ActivationRecord]]] = None
+    """
+    Random samples of activation records in each of the specified quantiles. None if quantile
+    tracking is disabled.
+    """
+    quantile_boundaries: Optional[list[float]] = None
+    """Boundaries of the quantiles used to generate the random_sample_by_quantile field."""
+
+    # Moments of activations
+    mean: Optional[float] = math.nan
+    variance: Optional[float] = math.nan
+    skewness: Optional[float] = math.nan
+    kurtosis: Optional[float] = math.nan
+
+    most_positive_activation_records: list[ActivationRecord] = field(
+        default_factory=list
+    )
+    """
+    Activation records with the most positive figure of merit value for this neuron over all dataset
+    examples.
+    """
+
+    @property
+    def max_activation(self) -> float:
+        """Return the maximum activation value over all top-activating activation records."""
+        return max(
+            [max(ar.activations) for ar in self.most_positive_activation_records]
+        )
+
+    def _get_top_activation_slices(
+        self, activation_record_slice_params: ActivationRecordSliceParams
+    ) -> dict[str, slice]:
+        splits = ["train", "calibration", "valid", "test"]
+        n_examples_per_split = activation_record_slice_params.n_examples_per_split
+        if n_examples_per_split is None:
+            n_examples_per_split = len(self.most_positive_activation_records) // len(
+                splits
+            )
+        assert len(self.most_positive_activation_records) >= n_examples_per_split * len(
+            splits
+        )
+        return get_slices_for_splits(splits, n_examples_per_split)
+
+    def _get_random_activation_slices(
+        self, activation_record_slice_params: ActivationRecordSliceParams
+    ) -> dict[str, slice]:
+        splits = ["calibration", "valid", "test"]
+        n_examples_per_split = activation_record_slice_params.n_examples_per_split
+        if n_examples_per_split is None:
+            n_examples_per_split = len(self.random_sample) // len(splits)
+        # NOTE: this assert could trigger on some old datasets with only 10 random samples, in which case you may have to remove "test" from the set of splits
+        assert len(self.random_sample) >= n_examples_per_split * len(splits)
+        return get_slices_for_splits(splits, n_examples_per_split)
+
+    def train_activation_records(
+        self,
+        activation_record_slice_params: ActivationRecordSliceParams,
+    ) -> list[ActivationRecord]:
+        """
+        Train split, typically used for generating explanations. Consists exclusively of
+        top-activating records since context window limitations make it difficult to include
+        random records.
+        """
+        return self.most_positive_activation_records[
+            self._get_top_activation_slices(activation_record_slice_params)["train"]
+        ]
+
+    def calibration_activation_records(
+        self,
+        activation_record_slice_params: ActivationRecordSliceParams,
+    ) -> list[ActivationRecord]:
+        """
+        Calibration split, typically used for calibrating neuron simulations. See
+        http://go/neuron_explanation_methodology for an explanation of calibration. Consists of
+        top-activating records and random records in a 1:1 ratio.
+        """
+        return (
+            self.most_positive_activation_records[
+                self._get_top_activation_slices(activation_record_slice_params)[
+                    "calibration"
+                ]
+            ]
+            + self.random_sample[
+                self._get_random_activation_slices(activation_record_slice_params)[
+                    "calibration"
+                ]
+            ]
+        )
+
+    def valid_activation_records(
+        self,
+        activation_record_slice_params: ActivationRecordSliceParams,
+    ) -> list[ActivationRecord]:
+        """
+        Validation split, typically used for evaluating explanations, either automatically with
+        simulation + correlation coefficient scoring, or manually by humans. Consists of
+        top-activating records and random records in a 1:1 ratio.
+        """
+        return (
+            self.most_positive_activation_records[
+                self._get_top_activation_slices(activation_record_slice_params)["valid"]
+            ]
+            + self.random_sample[
+                self._get_random_activation_slices(activation_record_slice_params)[
+                    "valid"
+                ]
+            ]
+        )
+
+    def test_activation_records(
+        self,
+        activation_record_slice_params: ActivationRecordSliceParams,
+    ) -> list[ActivationRecord]:
+        """
+        Test split, typically used for explanation evaluations that can't use the validation split.
+        Consists of top-activating records and random records in a 1:1 ratio.
+        """
+        return (
+            self.most_positive_activation_records[
+                self._get_top_activation_slices(activation_record_slice_params)["test"]
+            ]
+            + self.random_sample[
+                self._get_random_activation_slices(activation_record_slice_params)[
+                    "test"
+                ]
+            ]
+        )
+
+
+def neuron_exists(
+    dataset_path: str, layer_index: Union[str, int], neuron_index: Union[str, int]
+) -> bool:
+    """Return whether the specified neuron exists."""
+    file = bf.join(dataset_path, "neurons", str(layer_index), f"{neuron_index}.json")
+    return bf.exists(file)
+
+
+def load_neuron(
+    layer_index: Union[str, int],
+    neuron_index: Union[str, int],
+    dataset_path: str = "https://openaipublic.blob.core.windows.net/neuron-explainer/data/collated-activations",
+) -> NeuronRecord:
+    """Load the NeuronRecord for the specified neuron."""
+    url = "/".join([dataset_path, str(layer_index), f"{neuron_index}.json"])
+    url = standardize_azure_url(url)
+    with urllib.request.urlopen(url) as f:
+        neuron_record = loads(f.read())
+        if not isinstance(neuron_record, NeuronRecord):
+            raise ValueError(
+                f"Stored data incompatible with current version of NeuronRecord dataclass."
+            )
+        return neuron_record
+
+
+@bbb.ensure_session
+async def load_neuron_async(
+    layer_index: Union[str, int],
+    neuron_index: Union[str, int],
+    dataset_path: str = "az://openaipublic/neuron-explainer/data/collated-activations",
+) -> NeuronRecord:
+    """Async version of load_neuron."""
+    file = bf.join(dataset_path, str(layer_index), f"{neuron_index}.json")
+    return await read_neuron_file(file)
+
+
+@bbb.ensure_session
+async def read_neuron_file(neuron_filename: str) -> NeuronRecord:
+    """Like load_neuron_async, but takes a raw neuron filename."""
+    raw_contents = await bbb.read.read_single(neuron_filename)
+    neuron_record = loads(raw_contents.decode("utf-8"))
+    if not isinstance(neuron_record, NeuronRecord):
+        raise ValueError(
+            f"Stored data incompatible with current version of NeuronRecord dataclass."
+        )
+    return neuron_record
+
+
+def get_sorted_neuron_indices(
+    dataset_path: str, layer_index: Union[str, int]
+) -> List[int]:
+    """Returns the indices of all neurons in this layer, in ascending order."""
+    layer_dir = bf.join(dataset_path, "neurons", str(layer_index))
+    return sorted(
+        [
+            int(f.split(".")[0])
+            for f in bf.listdir(layer_dir)
+            if f.split(".")[0].isnumeric()
+        ]
+    )
+
+
+def get_sorted_layers(dataset_path: str) -> List[str]:
+    """
+    Return the indices of all layers in this dataset, in ascending numerical order, as strings.
+    """
+    return [
+        str(x)
+        for x in sorted(
+            [
+                int(x)
+                for x in bf.listdir(bf.join(dataset_path, "neurons"))
+                if x.isnumeric()
+            ]
+        )
+    ]

neuron_explainer/activations/attention_utils.py

@@ -0,0 +1,121 @@
+"""
+Contains math utilities for converting from flattened representations of attention activations
+(which are a scalar per token pair) to nested lists. The inner lists are attention activations
+related to attention from the same token (to different tokens).
+
+Tested in ./test_attention_utils.py.
+"""
+
+import math
+
+import numpy as np
+
+
+def _inverse_triangular_number(n: int) -> int:
+    # the m'th triangular number t_m satisfies t_m = m(m+1)/2
+    # this function asserts that n is a triangular number, and returns the unique m such that t_m = n
+    # this is used to infer the number of sequence tokens from the number of activations
+    assert n >= 0
+    m: int = (
+        math.floor(math.sqrt(1 + 8 * n)) - 1
+    ) // 2  # from quadratic formula applied to m(m+1)/2 = n
+    assert m * (m + 1) // 2 == n
+    return m
+
+
+def get_max_num_attended_to_sequence_tokens(num_sequence_tokens: int, num_activations: int) -> int:
+    # Attended to sequences are assumed to increase in length up to a maximum length, and then stay at that
+    # length for the remainder of the sequence. The maximum attended to sequence length is at most equal to the sequence length,
+    # but is permitted to be less
+    num_sequence_token_pairs = num_sequence_tokens * (num_sequence_tokens + 1) // 2
+    if num_activations == num_sequence_token_pairs:
+        # the maximum attended to sequence length is equal to the sequence length
+        return num_sequence_tokens
+    else:
+        # the maximum attended to sequence length is less than the sequence length, and
+        assert num_activations < num_sequence_token_pairs
+        num_missing_activations = num_sequence_token_pairs - num_activations
+        num_missing_sequence_tokens = _inverse_triangular_number(num_missing_activations)
+        max_num_attended_to_sequence_tokens = num_sequence_tokens - num_missing_sequence_tokens
+        assert max_num_attended_to_sequence_tokens > 0
+        return max_num_attended_to_sequence_tokens
+
+
+def get_attended_to_sequence_length_per_sequence_token(
+    num_sequence_tokens: int, max_num_attended_to_sequence_tokens: int
+) -> list[int]:
+    # given a num_sequence_tokens and a max_num_attended_to_sequence_tokens, return a list of length num_sequence_tokens
+    # where the ith element is the length of the attended to sequence for the ith sequence token.
+    # The length of the attended to sequence starts at 1, increases up to max_num_attended_to_sequence_tokens, by 1 with each
+    # token, and then stays at max_num_attended_to_sequence_tokens for the remainder of the sequence
+    assert num_sequence_tokens >= max_num_attended_to_sequence_tokens
+    attended_to_sequence_lengths = list(range(1, max_num_attended_to_sequence_tokens + 1))
+    if num_sequence_tokens > max_num_attended_to_sequence_tokens:
+        attended_to_sequence_lengths.extend(
+            [
+                max_num_attended_to_sequence_tokens
+                for _ in range(num_sequence_tokens - max_num_attended_to_sequence_tokens)
+            ]
+        )
+    return attended_to_sequence_lengths
+
+
+def get_attended_to_sequence_lengths(num_sequence_tokens: int, num_activations: int) -> list[int]:
+    max_num_attended_to_sequence_tokens = get_max_num_attended_to_sequence_tokens(
+        num_sequence_tokens, num_activations
+    )
+    return get_attended_to_sequence_length_per_sequence_token(
+        num_sequence_tokens, max_num_attended_to_sequence_tokens
+    )
+
+
+def _convert_flattened_index_to_unflattened_index_assuming_square_matrix(
+    flat_index: int,
+) -> tuple[int, int]:
+    # this con
+    n = math.floor((-1 + math.sqrt(1 + 8 * flat_index)) / 2)
+    m = flat_index - n * (n + 1) // 2
+    return n, m
+
+
+def convert_flattened_index_to_unflattened_index(
+    flattened_index: int,
+    num_sequence_tokens: int | None = None,
+    num_activations: int | None = None,
+) -> tuple[int, int]:
+    # given a flattened index, return the unflattened index
+    # if the attention matrix is square (most common), then the flattened_index uniquely determines the index within the square matrix
+    # if the attention matrix has more rows (sequence tokens) than columns (attended-to sequence tokens), then num_sequence_tokens
+    # and num_activations are required to determine the index within the matrix
+    # specify both num_sequence_tokens and num_activations, or neither
+    assert not (num_sequence_tokens is None) ^ (num_activations is None)
+
+    if (
+        num_sequence_tokens is None
+        or num_activations == num_sequence_tokens * (num_sequence_tokens + 1) // 2
+    ):
+        assume_square_matrix = True
+    else:
+        assume_square_matrix = False
+
+    if assume_square_matrix:
+        return _convert_flattened_index_to_unflattened_index_assuming_square_matrix(flattened_index)
+    else:
+        assert num_sequence_tokens is not None
+        assert num_activations is not None
+        assert flattened_index < num_activations
+        sequence_lengths = get_attended_to_sequence_lengths(num_sequence_tokens, num_activations)
+        sequence_lengths_cumsum = np.cumsum([0] + sequence_lengths)
+        sequence_index = int(
+            np.searchsorted(sequence_lengths_cumsum, flattened_index, side="right") - 1
+        )
+        assert sequence_lengths_cumsum[sequence_index] <= flattened_index, (
+            sequence_lengths_cumsum[sequence_index],
+            flattened_index,
+        )
+        assert sequence_lengths_cumsum[sequence_index + 1] >= flattened_index, (
+            sequence_lengths_cumsum[sequence_index + 1],
+            flattened_index,
+        )
+        index_within_sequence = flattened_index - sequence_lengths_cumsum[sequence_index]
+        return sequence_index, index_within_sequence