def sliced_transposed_product_concat(
    mat,
    block_size,
    axes=(-1,),
    precision=lax.Precision.DEFAULT,
):
    """Returns the concatenated slices representing mat*mat^T.

    Args:
      mat: The matrix for which we will compute mat*mat^T. It does not need to be
        square, and may be batched.
      block_size: The size of row blocks to compute.
      axes: Axes to use for the contraction.
      precision: The precision to use in each computation.

    Raises:
      ValueError: Raised when the specified block size does not evenly divide
        the number of rows of the input mat.
    """
    sliced_symmetric_matrix = sliced_transposed_product(
        mat=mat, block_size=block_size, axes=axes, precision=precision
    )
    return jnp.concatenate(sliced_symmetric_matrix.block_rows, axis=-1)
Returns the concatenated slices representing mat*mat^T.

    Args:
      mat: The matrix for which we will compute mat*mat^T. It does not need to be
        square, and may be batched.
      block_size: The size of row blocks to compute.
      axes: Axes to use for the contraction.
      precision: The precision to use in each computation.

    Raises:
      ValueError: Raised when the specified block size does not evenly divide
        the number of rows of the input mat.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py