def num_blocks_from_total_blocks(total_blocks):
    """Returns the number of blocks (i.e.

    block rows) from the total blocks.

    This is the inverse of the function x -> x*(x+1)/2.

    For example, the matrix M = [[A, B^T], [B, C]] may be represented using a
    total of 3 blocks ([A, B, C]). The number of corresponding block rows is 2.

    Args:
      total_blocks: The total blocks used to represent the matrix.
    """
    num_blocks = np.round((np.sqrt(8 * total_blocks + 1) - 1) / 2).astype(np.int32)
    if (num_blocks * (num_blocks + 1)) / 2 != total_blocks:
        raise ValueError(
            f"total_blocks={total_blocks} does not correspond to "
            "a symmetric matrix. It must have the form total_blocks = x*(x+1)/2."
        )
    return num_blocks
Returns the number of blocks (i.e.

    block rows) from the total blocks.

    This is the inverse of the function x -> x*(x+1)/2.

    For example, the matrix M = [[A, B^T], [B, C]] may be represented using a
    total of 3 blocks ([A, B, C]). The number of corresponding block rows is 2.

    Args:
      total_blocks: The total blocks used to represent the matrix.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py