def apply_logits(
        resid_post_mlp: torch.Tensor,
    ) -> torch.Tensor:
        """
        Input: (n_sequence_tokens, d_model) residual stream post-mlp activations at final layer.
        Output: (n_sequence_tokens, n_vocab) logprobs for each token in the sequence.
        """
        post_ln_f = apply_layer_norm(
            resid_post_mlp.unsqueeze(0),
            transformer.final_ln,
            detach_layer_norm_scale=detach_layer_norm_scale,
        )  # add batch dimension
        return transformer.unembed(post_ln_f).squeeze(0)  # remove batch dimension

        Input: (n_sequence_tokens, d_model) residual stream post-mlp activations at final layer.
        Output: (n_sequence_tokens, n_vocab) logprobs for each token in the sequence.
        
neuron_explainer/activations/derived_scalars/reconstituted.py