def slice_symmetric_matrix(
    mat,
    block_size,
):
    """Returns sliced row blocks.

    Args:
      mat: A symmetric matrix.
      block_size: The size of the row slices.
    """
    num_rows = mat.shape[-2]
    num_cols = mat.shape[-1]
    if num_rows != num_cols:
        raise ValueError("mat is not square.")
    if num_rows % block_size != 0:
        raise ValueError(
            f"block size does not evenly divide rows. num_rows={num_rows}, block_size={block_size}"
        )
    return SlicedSymmetricMatrix(
        block_rows=[
            mat[
                Ellipsis,
                i * block_size : (i + 1) * block_size,
                0 : (i + 1) * block_size,
            ]
            for i in range(num_rows // block_size)
        ]
    )
Returns sliced row blocks.

    Args:
      mat: A symmetric matrix.
      block_size: The size of the row slices.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py