def pad_block_symmetric_matrix(
    mat,
    symmetric_block_size,
    max_num_blocks,
):
    """Returns the padded blocked symmetric matrix.

    The size of the padded matrix will be:
      [symmetric_block_size, symmetric_block_size * max_num_blocks]

    The input matrix can either:
      - Be square with size less or equal to symmetric_block_size. In this case,
        mat will first be padded to a square matrix of size symmetric_block_size,
        and then be padded again up to the full size of the blocked matrix.
      - Be a rectangle with number of rows equal to block size.
        In this case, number of columns must be a multiple of number of rows, and
        the ratio must correspond to a block representation of a symmetric matrix.
        That is, the ratio must have form x * (x + 1) / 2. Here, x represents the
        number of block rows represented by the matrix.

    Args:
      mat: The input block matrix.
      symmetric_block_size: The size of blocks.
      max_num_blocks: The largest number of blocks to pad to.
    """
    rows, cols = mat.shape
    if rows > symmetric_block_size:
        raise ValueError(
            "Must have rows <= symmetric_block_size. Instead got "
            f"rows={rows}, symmetric_block_size={symmetric_block_size}."
        )
    if rows > cols:
        raise ValueError(
            f"Must have rows <= cols, instead got rows={rows}, cols={cols}."
        )
    if cols > symmetric_block_size * max_num_blocks:
        raise ValueError(
            "Must have cols <= symmetric_block_size * max_num_blocks "
            f"Instead got cols={cols}, "
            f"symmetric_block_size={symmetric_block_size}, "
            f"max_num_blocks={max_num_blocks}."
        )
    if rows < symmetric_block_size:
        mat = pad_square_matrix(mat, max_size=symmetric_block_size)
    # Update rows and cols after possibly padding in pad_square_matrix.
    rows, cols = mat.shape
    assert rows == symmetric_block_size
    assert cols % rows == 0
    filled_blocks = cols // rows
    padding_blocks = make_sliced_padding(
        symmetric_block_size=symmetric_block_size,
        num_blocks=symmetric_matrices.num_blocks_from_total_blocks(max_num_blocks),
        starting_block=symmetric_matrices.num_blocks_from_total_blocks(filled_blocks),
        dtype=mat.dtype,
    )
    return jnp.concatenate([mat, padding_blocks], axis=-1)
Returns the padded blocked symmetric matrix.

    The size of the padded matrix will be:
      [symmetric_block_size, symmetric_block_size * max_num_blocks]

    The input matrix can either:
      - Be square with size less or equal to symmetric_block_size. In this case,
        mat will first be padded to a square matrix of size symmetric_block_size,
        and then be padded again up to the full size of the blocked matrix.
      - Be a rectangle with number of rows equal to block size.
        In this case, number of columns must be a multiple of number of rows, and
        the ratio must correspond to a block representation of a symmetric matrix.
        That is, the ratio must have form x * (x + 1) / 2. Here, x represents the
        number of block rows represented by the matrix.

    Args:
      mat: The input block matrix.
      symmetric_block_size: The size of blocks.
      max_num_blocks: The largest number of blocks to pad to.
    
tools/train/scalable_shampoo/distributed_shampoo.py