def get_attn_write_tensor_by_layer_index(
    model_context: ModelContext,
    layer_indices: list[int] | None,
) -> dict[LayerIndex, torch.Tensor]:
    """Returns a dictionary mapping layer index to the write weight matrix for that layer."""
    if layer_indices is None:
        layer_indices = list(range(model_context.n_layers))
    W_out_by_layer_index: dict[LayerIndex, torch.Tensor] = {
        layer_index: model_context.get_weight(
            location_type=WeightLocationType.ATTN_TO_RESIDUAL,
            layer=layer_index,
            device=model_context.device,
        )  # shape (n_heads, d_head, d_model)
        for layer_index in layer_indices
    }
    return W_out_by_layer_index
Returns a dictionary mapping layer index to the write weight matrix for that layer.
neuron_explainer/activations/derived_scalars/write_tensors.py