def flatten_lower_triangle(t: torch.Tensor) -> torch.Tensor:
    """This turns the lower triangular entries of a tensor (that is, entries with
    i <= j for i and j indexing the last two dimensions) into a tensor flattened
    along the last two dimensions. This is used for turning attention-related activations
    which would normally be list[list[float]] for each token sequence, into a list[float],
    which is compatible with NeuronRecord's."""
    # Get the shape of the tensor
    shape = t.shape

    # Get the last dimension N (tensor should be ... x M x N)
    M, N = shape[-2:]

    # Get the indices of the lower triangular part
    row_indices, col_indices = torch.tril_indices(M, N)

    # Reshape the tensor so that all preceding dimensions are combined into one
    t_reshaped = t.reshape(-1, M, N)

    # Extract the lower triangular part and flatten it
    flattened_lower_triangle = t_reshaped[:, row_indices, col_indices]

    # Reshape back to separate out the original preceding dimensions
    new_shape = list(shape[:-2]) + [-1]
    flattened_lower_triangle = flattened_lower_triangle.reshape(new_shape)

    return flattened_lower_triangle
This turns the lower triangular entries of a tensor (that is, entries with
    i <= j for i and j indexing the last two dimensions) into a tensor flattened
    along the last two dimensions. This is used for turning attention-related activations
    which would normally be list[list[float]] for each token sequence, into a list[float],
    which is compatible with NeuronRecord's.
neuron_explainer/activations/derived_scalars/attention.py