def get_topk_of_tensor_list(
    tensor_list: list[torch.Tensor],
) -> tuple[torch.Tensor, list[int], list[tuple[int, ...]]]:
    """Given a list of tensors, returns the top k values and their indices (within the overall list and within each tensor)
    across all tensors in the list.
    """

    flattened_tensor = torch.cat([tensor.flatten() for tensor in tensor_list]).flatten()
    topk_values, topk_flat_indices = torch.topk(flattened_tensor, k=flattened_tensor.shape[0])
    topk_flat_indices = topk_flat_indices.cpu().numpy()
    cumsum_list_lengths = np.cumsum([tensor.numel() for tensor in tensor_list])

    def convert_flat_index_to_list_and_tensor_indices(
        flat_index: int,
    ) -> tuple[int, tuple[int, ...]]:
        list_index = np.searchsorted(cumsum_list_lengths, flat_index, side="right")
        if list_index > 0:
            flat_index -= cumsum_list_lengths[list_index - 1]
        tensor_indices = tuple(
            int(index) for index in np.unravel_index(flat_index, tensor_list[list_index].shape)
        )  # cast from np.int64 to int
        return (
            list_index,
            tensor_indices,  # type: ignore
        )

    list_indices_tuple, tensor_indices_tuple = zip(
        *[
            convert_flat_index_to_list_and_tensor_indices(flat_index)
            for flat_index in topk_flat_indices
        ]
    )

    return (
        topk_values,
        list(list_indices_tuple),
        list(tensor_indices_tuple),
    )
Given a list of tensors, returns the top k values and their indices (within the overall list and within each tensor)
    across all tensors in the list.
    
neuron_explainer/activations/derived_scalars/derived_scalar_store.py