def _compute_attn_write_norm_from_attn_and_value(
    attn_post_softmax: torch.Tensor,  # fth
    attn_value: torch.Tensor,  # thd
    W_O: torch.Tensor,  # hdo
    pass_type: PassType,
) -> torch.Tensor:
    """Computes the norm of the write vector from each token to each other token,
    by multiplying two saved activations and a model weight.
    For the forward pass, this is attn_post_softmax * ||W_O @ V|| (elementwise multiplication).
    For the backward pass (the gradient with respect to the forward pass quantity),
    this is grad_attn_post_softmax / ||W_O @ V|| (elementwise division)"""
    # attn_post_softmax: (nheads, num_sequence_tokens, num_attended_to_sequence_tokens)
    # attn_value: (num_attended_to_sequence_tokens, nheads, d_head)
    # W_O: (nheads, d_head, d_model)
    # output: (num_sequence_tokens, num_attended_to_sequence_tokens, nheads)
    num_sequence_tokens, num_attended_to_sequence_tokens, nheads = attn_post_softmax.shape
    d_head = attn_value.shape[-1]
    assert attn_value.shape == (num_attended_to_sequence_tokens, nheads, d_head)
    v_times_Wo_norm = _compute_v_times_Wo_norm(attn_value, W_O)
    assert v_times_Wo_norm.shape == (num_attended_to_sequence_tokens, nheads)
    if pass_type == PassType.FORWARD:
        v_times_Wo_norm_factor = v_times_Wo_norm
    else:
        assert pass_type == PassType.BACKWARD
        v_times_Wo_norm_factor = 1 / v_times_Wo_norm
    output = torch.einsum("fth,th->fth", attn_post_softmax, v_times_Wo_norm_factor)
    assert output.shape == (num_sequence_tokens, num_attended_to_sequence_tokens, nheads)
    return output
Computes the norm of the write vector from each token to each other token,
    by multiplying two saved activations and a model weight.
    For the forward pass, this is attn_post_softmax * ||W_O @ V|| (elementwise multiplication).
    For the backward pass (the gradient with respect to the forward pass quantity),
    this is grad_attn_post_softmax / ||W_O @ V|| (elementwise division)
neuron_explainer/activations/derived_scalars/attention.py