def row_abs_maxes(mat):
    """Returns the max of the absolute values of the rows of the full matrix.

    For example the symmetric matrix M = [[1, 6], [6, 2]] is represented using
    mat = [1, 6, 2] with block_size = 1. In this case the function returns the
    aboslute row maxes of the original symmetric matrix, [6, 6].

    Args:
      mat: The symmetric matrix represented as the concatenated blocks.
    """
    rows, cols = mat.shape

    # Find col and row max for each block.
    col_maxes = []
    row_maxes = []
    for i in range(cols // rows):
        block = jnp.abs(mat[Ellipsis, i * rows : (i + 1) * rows])
        col_maxes.append(jnp.max(block, axis=1))
        row_maxes.append(jnp.max(block, axis=0))

    # global row max from block maxes.
    num_blocks = num_blocks_from_total_blocks(cols // rows)
    maxes = []
    for i in range(num_blocks):
        maxes.append(
            jnp.concatenate(
                row_maxes[(i * (i + 1) // 2) : ((i + 2) * (i + 1) // 2)]
                + [
                    col_maxes[((j + 1) * (j + 2)) // 2 - (j - i + 1)]
                    for j in range(i + 1, num_blocks)
                ],
                axis=-1,
            )
        )

    return jnp.max(jnp.stack(maxes), axis=0)
Returns the max of the absolute values of the rows of the full matrix.

    For example the symmetric matrix M = [[1, 6], [6, 2]] is represented using
    mat = [1, 6, 2] with block_size = 1. In this case the function returns the
    aboslute row maxes of the original symmetric matrix, [6, 6].

    Args:
      mat: The symmetric matrix represented as the concatenated blocks.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py