def unflatten_lower_triangle_and_sum_columns(
    flattened: torch.Tensor, M: int, N: int
) -> torch.Tensor:
    """Equivalent to unflatten_lower_triangle(...).sum(dim=-1), less memory, more time.

    This function calculates the sum (over the last dimension) of the lower triangular part of a tensor
    from its flattened representation without instantiating the full matrix.
    """
    # Get the indices of the lower triangular part
    row_indices, col_indices = torch.tril_indices(M, N)
    num_elements = row_indices.shape[0]
    assert flattened.shape[-1] == num_elements

    # Create an empty tensor to store the sum of each row
    new_shape = list(flattened.shape[:-1]) + [M]
    reconstructed_summed = torch.zeros(new_shape, dtype=flattened.dtype, device=flattened.device)

    # Sum the elements over columns
    for i in range(num_elements):
        reconstructed_summed[..., row_indices[i]] += flattened[..., i]

    return reconstructed_summed
Equivalent to unflatten_lower_triangle(...).sum(dim=-1), less memory, more time.

    This function calculates the sum (over the last dimension) of the lower triangular part of a tensor
    from its flattened representation without instantiating the full matrix.
    
neuron_explainer/activations/derived_scalars/attention.py