def sliced_transposed_product(
    mat,
    block_size,
    axes=(-1,),
    precision=lax.Precision.DEFAULT,
):
    """Returns the blocked slices representing a symmetric contraction.

    Specifically, the output is a contraction of the input mat with itself, in the
    specified axes.

    Args:
      mat: The matrix for which we will compute a contraction with itself.
      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.
    """
    rank = len(mat.shape)

    def _make_axis_positive(ax):
        assert -rank <= ax < rank
        return ax + rank if ax < 0 else ax

    positive_axes = [_make_axis_positive(ax) for ax in axes]
    assert len(positive_axes) == len(axes)
    remaining_axes = set(range(rank)) - set(positive_axes)
    assert len(remaining_axes) == 1
    remaining_ax = remaining_axes.pop()

    num_rows = mat.shape[remaining_ax]
    if num_rows % block_size != 0:
        raise ValueError(
            "The row dimension must be divisible by block_size. "
            f"Instead got row dimension={num_rows} and block_size={block_size}."
        )

    block_rows = []
    for i in range(num_rows // block_size):
        start_indices = [0] * rank
        start_indices[remaining_ax] = i * block_size

        slice_sizes = list(mat.shape)
        slice_sizes[remaining_ax] = block_size

        slice_sizes_full = list(mat.shape)
        slice_sizes_full[remaining_ax] = (i + 1) * block_size

        block_rows.append(
            product_with_transpose(
                lax.dynamic_slice(
                    mat, start_indices=start_indices, slice_sizes=slice_sizes
                ),
                lax.dynamic_slice(
                    mat, start_indices=[0] * rank, slice_sizes=slice_sizes_full
                ),
                axes=(axes, axes),
                precision=precision,
            )
        )

    return SlicedSymmetricMatrix(block_rows=block_rows)
Returns the blocked slices representing a symmetric contraction.

    Specifically, the output is a contraction of the input mat with itself, in the
    specified axes.

    Args:
      mat: The matrix for which we will compute a contraction with itself.
      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