def maybe_construct_transform_fns_for_top_k(
    token_index: int | None,
) -> tuple[Callable[[torch.Tensor], torch.Tensor] | None, Callable[[NodeIndex], NodeIndex] | None]:
    """
    In some cases, we may want to compute the top k activations on a transformed version of the DerivedScalarStore.
    This function constructs a transformation to apply to the activations, and a corresponding transformation to
    apply to the NodeIndex objects returned by the .topk operation, if any, to index the un-transformed DerivedScalarStore.
    `None` means no transformation is to be performed.
    Currently, we use this for computing top k over a single token index only, and converting the NodeIndex objects
    to index the original, un-transformed DerivedScalarStore.
    """
    if token_index is None:
        return None, None
    else:

        def access_token_index(x: torch.Tensor) -> torch.Tensor:
            return x[token_index].unsqueeze(0)

        def convert_to_original_token_index(node_index: NodeIndex) -> NodeIndex:
            assert (
                node_index.tensor_indices[0] == 0
            )  # assumed that token dimension was converted to singleton
            # replace the 0 token index with the original token_index
            return node_index.with_updates(
                tensor_indices=(token_index, *node_index.tensor_indices[1:])
            )

        return access_token_index, convert_to_original_token_index

    In some cases, we may want to compute the top k activations on a transformed version of the DerivedScalarStore.
    This function constructs a transformation to apply to the activations, and a corresponding transformation to
    apply to the NodeIndex objects returned by the .topk operation, if any, to index the un-transformed DerivedScalarStore.
    `None` means no transformation is to be performed.
    Currently, we use this for computing top k over a single token index only, and converting the NodeIndex objects
    to index the original, un-transformed DerivedScalarStore.
    
neuron_explainer/activation_server/interactive_model.py