def topk(
        self,
        k: int | None,
        pass_type: PassType = PassType.FORWARD,
        dsts: list[DerivedScalarType] | None = None,
        layer_indices: list[int] | None = None,
        largest: bool = True,
    ) -> tuple[torch.Tensor, list[DerivedScalarIndex]]:
        """The strategy here, similar to the topk method of ActivationsAndMetadata, is to compute
        topk for each DST separately, and then combine them. This avoids instantiating a big tensor
        containing the entire contents of the DerivedScalarStore at any point."""
        preprocessed = self
        if layer_indices is not None:
            preprocessed = preprocessed.filter_layers(layer_indices)
        if dsts is not None:
            preprocessed = preprocessed.filter_dsts(dsts)
        preprocessed = preprocessed.filter_pass_type(pass_type)
        if k is None:
            k = preprocessed.numel()
        assert k is not None

        topk_by_dst_and_pass_type = {
            (dst, pass_type): self.activations_and_metadata_by_dst_and_pass_type[
                (dst, pass_type)
            ].topk(
                k=k,
                largest=largest,
            )
            for dst in preprocessed.dsts
        }

        (
            overall_topk_values,
            overall_topk_list_indices,
            overall_topk_tensor_indices,
        ) = get_topk_of_tensor_list([topk[0] for topk in topk_by_dst_and_pass_type.values()])

        topk_ds_indices = [topk[1] for topk in topk_by_dst_and_pass_type.values()]

        overall_topk_ds_indices = []
        for list_index, tensor_indices in zip(
            overall_topk_list_indices, overall_topk_tensor_indices
        ):
            assert len(tensor_indices) == 1
            overall_topk_ds_indices.append(topk_ds_indices[list_index][tensor_indices[0]])

        return overall_topk_values, overall_topk_ds_indices
The strategy here, similar to the topk method of ActivationsAndMetadata, is to compute
        topk for each DST separately, and then combine them. This avoids instantiating a big tensor
        containing the entire contents of the DerivedScalarStore at any point.
neuron_explainer/activations/derived_scalars/derived_scalar_store.py