def materialize_matrix(symmetric_matrix):
    """Returns a materialized symmetric matrix.

    Args:
      symmetric_matrix: the matrix represented by lower-triangular block slices.
    """
    block_rows = symmetric_matrix.block_rows
    block_size = block_rows[0].shape[-2]
    num_blocks = len(block_rows)

    # Slice the lower-triangular and diagonal blocks into blocks.
    blocks = [
        [
            block_row[Ellipsis, i * block_size : (i + 1) * block_size]
            for i in range(k + 1)
        ]
        for k, block_row in enumerate(block_rows)
    ]

    # Generate the (off-diagonal) upper-triangular blocks.
    off_diags = [[] for _ in range(num_blocks - 1)]
    for k, block_row in enumerate(block_rows[1:]):
        for i in range(k + 1):
            off_diags[i].append(
                jnp.swapaxes(
                    a=block_row[Ellipsis, i * block_size : (i + 1) * block_size],
                    axis1=-1,
                    axis2=-2,
                )
            )

    return jnp.block(
        [row + row_t for row, row_t in zip(blocks[:-1], off_diags)] + [blocks[-1]]
    )
Returns a materialized symmetric matrix.

    Args:
      symmetric_matrix: the matrix represented by lower-triangular block slices.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py