def _get_token_pair_attribution(
        self,
        request_spec: TokenPairAttributionRequestSpec,
        multi_group_ds_store: MultiGroupDerivedScalarStore,
        all_node_indices: list[NodeIndex],
    ) -> TokenPairAttributionResponseData:
        """Returns attended-to tokens with most positive attributions for attention write autoencoder latents."""
        ds_store = multi_group_ds_store.get_ds_store(group_id=GroupId.TOKEN_PAIR_ATTRIBUTION)
        postprocessor = TokenPairAttributionConverter(
            model_context=self._standard_model_context,
            multi_autoencoder_context=self._multi_autoencoder_context,
            num_tokens_attended_to=request_spec.num_tokens_attended_to,
        )

        # sort the top token-attended-to by the value of the attribution
        node_indices = []  # type: list[MirroredNodeIndex]
        top_tokens_attended_to_list = []  # type: list[TopTokensAttendedTo | None]
        for node_index in all_node_indices:
            node_indices.append(MirroredNodeIndex.from_node_index(node_index))
            try:
                postprocessed = postprocessor.postprocess(node_index, ds_store)
                top_tokens_attended_to = postprocessed.topk(k=request_spec.num_tokens_attended_to)
                top_tokens_attended_to_list.append(
                    TopTokensAttendedTo(
                        token_indices=top_tokens_attended_to.indices.cpu().numpy().tolist(),
                        attributions=top_tokens_attended_to.values.cpu().numpy().tolist(),
                    )
                )
            except ValueError:
                top_tokens_attended_to_list.append(None)
                continue

        return TokenPairAttributionResponseData(
            node_indices=node_indices, top_tokens_attended_to_list=top_tokens_attended_to_list
        )
Returns attended-to tokens with most positive attributions for attention write autoencoder latents.
neuron_explainer/activation_server/interactive_model.py