def slice_symmetric_matrix_concat(
    mat,
    block_size,
):
    """Returns the concatenated sliced row blocks.

    Args:
      mat: A symmetric matrix.
      block_size: The size of the row slices.
    """
    sliced_symmetric_matrix = slice_symmetric_matrix(mat=mat, block_size=block_size)
    return jnp.concatenate(sliced_symmetric_matrix.block_rows, axis=-1)
Returns the concatenated 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