def make_vocab_token_write_to_input_direction_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """
    This computes the inner product between each token's embedding vector, and the input
    weight for a particular autoencoder latent or MLP neuron. The MLP neuron case
    straightforwardly applies the logit-lens setup: Embedding * InputWeightWithLayerNormGain.
    For the MLP autoencoder case (so far only implemented for MLP latents), we first convert
    to the MLP neuron basis using MlpLayer(Embedding), then multiply by InputWeight (there is
    no associated layer norm gain). MlpLayer(Embedding) * InputWeight corresponds to the
    autoencoder latent pre-activation in the case where the residual stream vector were exactly
    the token embedding vector.
    """
    dst = DerivedScalarType.VOCAB_TOKEN_WRITE_TO_INPUT_DIRECTION

    activation_index = dst_config.activation_index_for_fake_grad
    assert activation_index is not None
    layer_index = activation_index.layer_index
    assert layer_index is not None
    activation_location_type = activation_index.activation_location_type
    neuron_index = activation_index.tensor_indices[1]
    assert isinstance(neuron_index, int)

    model_context = dst_config.get_model_context()
    autoencoder_context = dst_config.get_autoencoder_context(NodeType.MLP_AUTOENCODER_LATENT)

    assert isinstance(model_context, StandardModelContext)
    transformer = model_context.get_or_create_model()
    transformer_layer = transformer.xf_layers[layer_index]

    assert activation_location_type in (
        ActivationLocationType.MLP_POST_ACT,
        ActivationLocationType.ONLINE_AUTOENCODER_LATENT,
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_LATENT,
    ), f"Virtual input weight not defined for {activation_location_type=}"

    emb = get_embedding(model_context)  # resulting shape: n_vocab, d_model
    match activation_location_type:
        # only fwd pass is used, so detaching layer norm scale is not relevant
        case ActivationLocationType.MLP_POST_ACT:
            token_vector = apply_mlp_pre_act(
                transformer_layer, emb, detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE
            )[:, neuron_index]
            # ^ resulting shape: n_vocab,
        case (
            ActivationLocationType.ONLINE_MLP_AUTOENCODER_LATENT
            | ActivationLocationType.ONLINE_AUTOENCODER_LATENT
        ):
            assert autoencoder_context is not None
            autoencoder = autoencoder_context.get_autoencoder(layer_index)
            # directly applies MLP layer + autoencoder to all token embedding
            # vectors
            token_vector = apply_autoencoder_pre_latent(
                transformer_layer,
                autoencoder,
                emb,
                autoencoder_dst=autoencoder_context.dst,
                detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE,
                # index the latent to avoid creating a matrix of shape (n_vocab, n_latents)
                latent_slice=slice(neuron_index, neuron_index + 1),
            )[:, 0]
            # indexed by zero because we already sliced to a single neuron
            # resulting shape: n_vocab,
        case _:
            raise NotImplementedError(activation_location_type)

    def replace_with_token_vector(
        resid_post_emb: torch.Tensor,
    ) -> torch.Tensor:
        # resid_post_emb only used for shape
        return token_vector.unsqueeze(0).expand(  # resulting shape: 1, n_vocab
            resid_post_emb.shape[0], -1
        )  # resulting shape: n_sequence_tokens, n_vocab

    resid_post_emb_scalar_deriver = make_scalar_deriver_factory_for_activation_location_type(
        ActivationLocationType.RESID_POST_EMBEDDING
    )(dst_config)

    return resid_post_emb_scalar_deriver.apply_transform_fn_to_output(
        replace_with_token_vector, pass_type_to_transform=PassType.FORWARD, output_dst=dst
    )

    This computes the inner product between each token's embedding vector, and the input
    weight for a particular autoencoder latent or MLP neuron. The MLP neuron case
    straightforwardly applies the logit-lens setup: Embedding * InputWeightWithLayerNormGain.
    For the MLP autoencoder case (so far only implemented for MLP latents), we first convert
    to the MLP neuron basis using MlpLayer(Embedding), then multiply by InputWeight (there is
    no associated layer norm gain). MlpLayer(Embedding) * InputWeight corresponds to the
    autoencoder latent pre-activation in the case where the residual stream vector were exactly
    the token embedding vector.
    
neuron_explainer/activations/derived_scalars/residual.py