def find_num_blocks(block_rows_concat):
    """Returns the number of (row) blocks representing the concatenated matrix.

    For example, an input with dimensions [256, 2560] represents 10 square blocks,
    which matches 4 lower-triangular block rows (1+2+3+4). So this function will
    return 4.

    Use ordinary numpy functions here so that the returned value is static.

    Args:
      block_rows_concat: The concatenated block array.

    Raises:
      ValueError: When the dimensions of the matrix do not correspond to a lower
      triangular block representation.
    """
    # Compute the number of square blocks used to represent the matrix.
    total_blocks = block_rows_concat.shape[-1] / block_rows_concat.shape[-2]
    # Determine the number of block rows by inverting y = x*(x+1)/2.
    return num_blocks_from_total_blocks(total_blocks)
Returns the number of (row) blocks representing the concatenated matrix.

    For example, an input with dimensions [256, 2560] represents 10 square blocks,
    which matches 4 lower-triangular block rows (1+2+3+4). So this function will
    return 4.

    Use ordinary numpy functions here so that the returned value is static.

    Args:
      block_rows_concat: The concatenated block array.

    Raises:
      ValueError: When the dimensions of the matrix do not correspond to a lower
      triangular block representation.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py