def sliced_matrix_diag(mat):
    """Returns the diagonal of the symmetric matrix.

    Args:
      mat: The symmetric matrix represented in concatenated block form.
    """
    rows, cols = mat.shape
    total_blocks = cols // rows
    num_blocks = num_blocks_from_total_blocks(total_blocks)
    diags = []
    for i in range(num_blocks):
        last_index = rows * ((i + 2) * (i + 1)) // 2
        first_index = last_index - rows
        diags.append(jnp.diag(mat[Ellipsis, first_index:last_index]))
    return jnp.concatenate(diags, axis=-1)
Returns the diagonal of the symmetric matrix.

    Args:
      mat: The symmetric matrix represented in concatenated block form.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py