def times_vector(mat, vec):
    """Returns the symmetric block-concatenated matrix multiplied by a vector.

    Specifically, each value in the vector is multiplied by a row of the full
    matrix. That is, the vector is broadcast and multiplied element-wise. Note
    this would be the transpose of full_mat * vec if full_mat represented the full
    symmetric matrix.

    Args:
      mat: The symmetric matrix represented as the concatenated blocks.
      vec: The vector, having the same dimension as the materialized matrix.
    """
    rows, cols = mat.shape
    num_blocks = num_blocks_from_total_blocks(cols // rows)
    multiplied = []
    for i in range(num_blocks):
        mat_block = mat[
            Ellipsis, rows * ((i + 1) * i) // 2 : rows * ((i + 1) * (i + 2)) // 2
        ]
        vec_block = vec[Ellipsis, rows * i : rows * (i + 1)]
        multiplied.append(jnp.einsum("...ij,...i->ij", mat_block, vec_block))
    return jnp.concatenate(multiplied, axis=-1)
Returns the symmetric block-concatenated matrix multiplied by a vector.

    Specifically, each value in the vector is multiplied by a row of the full
    matrix. That is, the vector is broadcast and multiplied element-wise. Note
    this would be the transpose of full_mat * vec if full_mat represented the full
    symmetric matrix.

    Args:
      mat: The symmetric matrix represented as the concatenated blocks.
      vec: The vector, having the same dimension as the materialized matrix.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py