def get_embedding(model_context: ModelContext) -> torch.Tensor:
    """
    returns an embedding matrix. Note that there is no layer norm in between the embedding tensor and
    the residual stream
    """
    return model_context.get_weight(
        location_type=WeightLocationType.EMBEDDING,
        device=model_context.device,
    )

    returns an embedding matrix. Note that there is no layer norm in between the embedding tensor and
    the residual stream
    
neuron_explainer/models/model_context.py