def unflatten_lower_triangle(flattened: torch.Tensor, M: int, N: int) -> torch.Tensor:
    """Inverse of flatten_lower_triangle.

    This function reconstructs a tensor from its flattened lower triangular part.
    The upper triangular part of the reconstructed tensor is filled with zeros.
    """
    # Get the indices of the lower triangular part
    row_indices, col_indices = torch.tril_indices(M, N)

    # Create an empty tensor to fill in the original shape
    new_shape = list(flattened.shape[:-1]) + [M, N]
    reconstructed = torch.zeros(new_shape, dtype=flattened.dtype, device=flattened.device)

    # Fill in the lower triangular part of the reconstructed tensor
    reconstructed[..., row_indices, col_indices] = flattened

    return reconstructed
Inverse of flatten_lower_triangle.

    This function reconstructs a tensor from its flattened lower triangular part.
    The upper triangular part of the reconstructed tensor is filled with zeros.
    
neuron_explainer/activations/derived_scalars/attention.py