def make_sliced_padding(
    symmetric_block_size,
    num_blocks,
    starting_block,
    dtype,
):
    """Returns padding for symmetric block matrix.

    Specifically, the padding is given concatenated rectangular matrices
    representing the lower-triangular rows below the starting block. For example,
    if we want to pad the symmetric matrix

    M = [[A, B^T]
         [B, C]],

    the desired output (in terms of the full matrix) with num_blocks = 4 is

    M_padded = [[A, B^T, 0, 0]
                [B, C,   0, 0]
                [0, 0,   I, 0]
                 0, 0,   0, I].

    We would represent M as the block matrix mat = [A, B, C]. In this form, the
    additional padding to provide has form [0, 0, I, 0, 0, 0, I] (only the lower
    triangular parts in the third and fourth rows).

    Args:
      symmetric_block_size: The size of each block.
      num_blocks: The total number of blocks.
      starting_block: The block where to start the padding.
      dtype: The type to use for the blocks.
    """
    if starting_block == num_blocks:
        return jnp.zeros(shape=(symmetric_block_size, 0), dtype=dtype)

    blocks = []
    for i in range(starting_block, num_blocks):
        blocks.append(
            jnp.zeros(
                shape=(symmetric_block_size, symmetric_block_size * i), dtype=dtype
            )
        )
        blocks.append(jnp.eye(symmetric_block_size, dtype=dtype))
    return jnp.concatenate(blocks, axis=-1)
Returns padding for symmetric block matrix.

    Specifically, the padding is given concatenated rectangular matrices
    representing the lower-triangular rows below the starting block. For example,
    if we want to pad the symmetric matrix

    M = [[A, B^T]
         [B, C]],

    the desired output (in terms of the full matrix) with num_blocks = 4 is

    M_padded = [[A, B^T, 0, 0]
                [B, C,   0, 0]
                [0, 0,   I, 0]
                 0, 0,   0, I].

    We would represent M as the block matrix mat = [A, B, C]. In this form, the
    additional padding to provide has form [0, 0, I, 0, 0, 0, I] (only the lower
    triangular parts in the third and fourth rows).

    Args:
      symmetric_block_size: The size of each block.
      num_blocks: The total number of blocks.
      starting_block: The block where to start the padding.
      dtype: The type to use for the blocks.
    
tools/train/scalable_shampoo/distributed_shampoo.py