def _convert_rectangular_activations_to_flattened_lower_triangle(
    rectangular_activations: torch.Tensor,
) -> torch.Tensor:
    """
    This function takes a tensor with the first two dimensions indexed by rows i and columns j.
    It keeps the entries where i <= j and flattens the first 2 dimensions into a 1-d list.
    """

    assert rectangular_activations.ndim == 3, rectangular_activations.shape
    num_sequence_tokens, num_attended_to_sequence_tokens, nheads = rectangular_activations.shape
    rectangular_activations = torch.einsum(
        "fth->hft", rectangular_activations
    )  # (nheads, num_sequence_tokens, num_attended_to_sequence_tokens
    flattened_activations = flatten_lower_triangle(rectangular_activations)
    flattened_activations = torch.einsum(
        "hs->sh", flattened_activations
    )  # (num_token_pairs, nheads)
    min_num_tokens = min(num_sequence_tokens, num_attended_to_sequence_tokens)
    num_token_pairs_before_attended_to_sequence_is_saturated = (
        min_num_tokens * (min_num_tokens + 1) // 2
    )
    if num_sequence_tokens > num_attended_to_sequence_tokens:
        num_token_pairs_after_attended_to_sequence_is_saturated = (
            num_sequence_tokens - num_attended_to_sequence_tokens
        ) * num_attended_to_sequence_tokens
        num_token_pairs = (
            num_token_pairs_before_attended_to_sequence_is_saturated
            + num_token_pairs_after_attended_to_sequence_is_saturated
        )
    else:
        num_token_pairs = num_token_pairs_before_attended_to_sequence_is_saturated
    # if num_sequence_tokens == num_attended_to_sequence_tokens, then the i>=j entries form a lower triangle (first term)
    # if num_sequence_tokens > num_attended_to_sequence_tokens, then there is an additional rectangle of entries
    # below that lower triangle with i > j (second term)
    assert flattened_activations.shape == (
        num_token_pairs,
        nheads,
    ), f"{flattened_activations.shape=} != {(num_token_pairs, nheads)=}"
    return flattened_activations

    This function takes a tensor with the first two dimensions indexed by rows i and columns j.
    It keeps the entries where i <= j and flattens the first 2 dimensions into a 1-d list.
    
neuron_explainer/activations/derived_scalars/attention.py