tf.einsum是否依赖输入矩阵尺寸？256x256矩阵场景异常问询

问题背景 #

我使用以下代码从分解后的Cholesky矩阵构造密度矩阵：

def density_matrix_from_T(tmatrix):
    """
    Gets density matrices from T matrices and normalizes them.

    Args:
        tmatrix (`tf.Tensor`): A tensor (N, hilbert_size, hilbert_size)
                               representing N valid T matrices.

    Returns:
        rho (`tf.Tensor`): A tensor of shape (N, hilbert_size, hilbert_size)
                           representing N density matrices.
    """

    T = tmatrix
    T_dagger = tf.transpose(T, perm=[0, 2, 1], conjugate=True)
    proper_dm = tf.matmul(T_dagger, T)
    all_traces = tf.linalg.trace(proper_dm)
    all_traces = tf.reshape(1 / all_traces, (-1, 1))
    rho = tf.einsum("bij,bk->bij", proper_dm, all_traces)

    return rho

我的疑问集中在rho = tf.einsum("bij,bk->bij", proper_dm, all_traces)这一行：当输入为16x16矩阵时，代码运行良好，生成的rho矩阵与参考矩阵相似度达0.99；但输入为256x256矩阵时，相似度仅为0.005，结果很差。我不确定这段代码是否输出了预期结果，想咨询：

• tf.einsum是否会以某种方式依赖输入矩阵的尺寸？
• 该用法对于256x256矩阵是否正确？

解答 #

关于tf.einsum的尺寸依赖 #

tf.einsum本身不依赖输入矩阵的尺寸，它严格按照你指定的下标规则执行张量运算（收缩、广播等），只要输入张量的下标匹配规则，不管维度是16还是256，逻辑都是一致的。

你的einsum用法是否正确 #

你写的"bij,bk->bij"是合法且逻辑正确的：

• proper_dm是形状为(N, H, H)的张量（H为希尔伯特空间维度）
• all_traces被reshape为(N, 1)，einsum中k维度长度为1，会自动广播到proper_dm的i、j维度
• 本质是对每个batch内的H×H矩阵，每个元素都除以该batch的迹，这完全符合密度矩阵的归一化要求。

256x256矩阵结果差的可能原因 #

问题不在einsum，而是其他环节：

1. 数值精度问题：256x256矩阵的迹计算容易出现数值溢出/下溢。比如T矩阵元素量级失控时，T_dagger @ T的迹会变得极大或极小，导致归一化后数值失真。建议检查all_traces的取值范围，看是否存在无穷大、NaN或极端值。
2. T矩阵合法性：Cholesky分解的结果必须是下三角矩阵（且对角线元素为正），如果256x256的T矩阵不符合这个结构，T_dagger @ T就不是半正定矩阵，无法得到合法的密度矩阵，自然和参考矩阵相似度极低。
3. 相似度计算方式：高维度矩阵用普通元素相似度（如余弦相似度）会因维度稀释导致数值偏低，应该使用密度矩阵专用的相似度指标，比如保真度（Fidelity）。

代码简化建议 #

你的einsum逻辑可以用更直观的广播除法替代，可读性更高，且逻辑完全一致：

rho = proper_dm / tf.expand_dims(all_traces, axis=[1, 2])

替换后可排除einsum的潜在误解，再测试256x256的情况。

内容的提问来源于stack exchange，提问作者Silvio sjsj