def product_with_transpose(
    mat1,
    mat2,
    axes,
    precision=lax.Precision.DEFAULT,
):
    """Returns mat1 * mat2^T for two matrices (possibly batched).

    The rows and columns are the last two dimensions for each matrix.

    Args:
      mat1: First matrix.
      mat2: Second matrix.
      axes: The axes over which to apply the product.
      precision: JAX precision to use for the multiplication.
    """
    return jnp.tensordot(a=mat1, b=mat2, axes=axes, precision=precision)
Returns mat1 * mat2^T for two matrices (possibly batched).

    The rows and columns are the last two dimensions for each matrix.

    Args:
      mat1: First matrix.
      mat2: Second matrix.
      axes: The axes over which to apply the product.
      precision: JAX precision to use for the multiplication.
    
tools/train/scalable_shampoo/symmetric_matrices/symmetric_matrices.py