def get_unembedding_with_ln_gain(model_context: ModelContext) -> torch.Tensor:
    """
    returns an unembedding matrix multiplied by the layer norm gain (a d_model-dimensional vector)
    for the final layer
    """
    Unemb_without_ln_gain = model_context.get_weight(
        location_type=WeightLocationType.UNEMBEDDING,
        device=model_context.device,
    )
    ln_gain_final = model_context.get_weight(
        location_type=WeightLocationType.LAYER_NORM_GAIN_FINAL,
        device=model_context.device,
    )
    return torch.einsum("ov,o->ov", Unemb_without_ln_gain, ln_gain_final)

    returns an unembedding matrix multiplied by the layer norm gain (a d_model-dimensional vector)
    for the final layer
    
neuron_explainer/models/model_context.py