def _compute_v_times_Wo_norm(
    attn_value: torch.Tensor,  # thd
    W_O: torch.Tensor,  # hdo
) -> torch.Tensor:
    """
    This function computes the norm of the product of attention values and the output weight matrix
    (W_O in the transformer-circuits convention, or c_proj in many internal settings).
    The attention values are represented by the tensor attn_value and the output weight matrix by W_O.
    The function returns a tensor representing the norm of the product.
    """

    num_attended_to_sequence_tokens, nheads, d_head = attn_value.shape
    assert W_O.shape[:2] == (nheads, d_head)  # third dim is d_model
    v_times_Wo = torch.einsum("thd,hdo->tho", attn_value, W_O)
    return torch.linalg.norm(v_times_Wo, dim=-1)

    This function computes the norm of the product of attention values and the output weight matrix
    (W_O in the transformer-circuits convention, or c_proj in many internal settings).
    The attention values are represented by the tensor attn_value and the output weight matrix by W_O.
    The function returns a tensor representing the norm of the product.
    
neuron_explainer/activations/derived_scalars/attention.py