def test_unflatten_lower_triangle(extra_dim: list[int], N: int) -> None:
    """Test that unflatten_lower_triangle is the inverse of flatten_lower_triangle."""
    # Create a random tensor of shape ... x M x N
    M = 64
    original_tensor = torch.rand(extra_dim + [M, N])

    # Set all elements above the lower triangular to 0
    lower_triangular_mask = torch.tril(torch.ones(M, N)).bool()
    original_tensor[..., ~lower_triangular_mask] = 0

    # Apply flatten_lower_triangle to the original tensor
    flattened = flatten_lower_triangle(original_tensor)
    assert flattened.shape == tuple(extra_dim + [lower_triangular_mask.sum()])

    # Apply unflatten_lower_triangle to the flattened tensor
    reconstructed_tensor = unflatten_lower_triangle(flattened, M, N)
    assert torch.allclose(original_tensor, reconstructed_tensor)
Test that unflatten_lower_triangle is the inverse of flatten_lower_triangle.
neuron_explainer/activations/derived_scalars/tests/test_attention.py