def update_sliced_rows(
    symmetric_matrix,
    mat,
    alpha,
    beta,
    axes=(-1,),
):
    """Implements the blocked equivalent of SYRK.

    Specifically, the symmetric matrix (represented using lower-triangular block
    rows) is updated using the sliced product of mat.

    Args:
      symmetric_matrix: The symmetric matrix to update.
      mat: The matrix to use for the update = mat * mat^T. The number of rows
        should match that of symmetric_matrix.
      alpha: The weight for the update.
      beta: The weight for the original symmetric matrix.
      axes: Axes to use for the contraction of the update.

    Returns:
      The updated rows of alpha * mat * mat^T + beta * symmetric_matrix.
    """
    block_size = symmetric_matrix.block_rows[0].shape[-2]
    sym_prod = sliced_transposed_product(mat=mat, block_size=block_size, axes=axes)
    return SlicedSymmetricMatrix(
        block_rows=[
            update * alpha + row * beta
            for update, row in zip(sym_prod.block_rows, symmetric_matrix.block_rows)
        ]
    )
Implements the blocked equivalent of SYRK.

    Specifically, the symmetric matrix (represented using lower-triangular block
    rows) is updated using the sliced product of mat.

    Args:
      symmetric_matrix: The symmetric matrix to update.
      mat: The matrix to use for the update = mat * mat^T. The number of rows
        should match that of symmetric_matrix.
      alpha: The weight for the update.
      beta: The weight for the original symmetric matrix.
      axes: Axes to use for the contraction of the update.

    Returns:
      The updated rows of alpha * mat * mat^T + beta * symmetric_matrix.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py