def _get_attn_write_for_one_layer_index(
    model_context: ModelContext,
    layer_index: int,
) -> torch.Tensor:
    """Returns a dictionary mapping layer index to the write weight matrix for that layer."""
    return get_attn_write_tensor_by_layer_index(
        model_context=model_context,
        layer_indices=[layer_index],
    )[layer_index]
Returns a dictionary mapping layer index to the write weight matrix for that layer.
neuron_explainer/activations/derived_scalars/attention.py