def materialize_matrix_from_concat(
    block_rows_concat,
    num_blocks=None,
):
    """Returns a materialized symmetric matrix from concatenated slices.

    Args:
      block_rows_concat: The matrix represented as the concatenated
        lower-triangular blocks.
      num_blocks: The number of block-rows used to represent the symmetric matrix.
        If not specified, it is inferred from the shape of block_rows_concat.
    """
    if num_blocks is None:
        num_blocks = find_num_blocks(block_rows_concat)

    block_size = block_rows_concat.shape[-2]

    block_rows = [
        block_rows_concat[
            Ellipsis,
            (k * (k + 1))
            // 2
            * block_size : (((k + 1) * (k + 2)) // 2 + 1)
            * block_size,
        ]
        for k in range(num_blocks)
    ]

    return materialize_matrix(SlicedSymmetricMatrix(block_rows=block_rows))
Returns a materialized symmetric matrix from concatenated slices.

    Args:
      block_rows_concat: The matrix represented as the concatenated
        lower-triangular blocks.
      num_blocks: The number of block-rows used to represent the symmetric matrix.
        If not specified, it is inferred from the shape of block_rows_concat.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py