def diag_as_concat(diag, block_size):
    """Returns the representation of a diagonal matrix in symmetric block form.

    Args:
      diag: The 1D array for the diagonals.
      block_size: The size of blocks to use. Must divide the length of diag.
    """
    assert len(diag.shape) == 1  # diag must be 1D.
    assert len(diag) % block_size == 0
    num_diag_blocks = len(diag) // block_size
    blocks = []
    for i in range(num_diag_blocks):
        blocks.append(jnp.zeros(shape=(block_size, block_size * i), dtype=diag.dtype))
        blocks.append(jnp.diag(diag[i * block_size : (i + 1) * block_size]))
    return jnp.concatenate(blocks, axis=-1)
Returns the representation of a diagonal matrix in symmetric block form.

    Args:
      diag: The 1D array for the diagonals.
      block_size: The size of blocks to use. Must divide the length of diag.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py