TensorFlow实现：找出行内元素大于对并堆叠成n×m×m张量

核心实现方案 #

直接利用TensorFlow的广播机制完成向量化操作，无需循环，高效生成目标张量：

import tensorflow as tf

# 示例输入张量
X = tf.constant([[1,2],[4,3], [5,7]])
# 通过维度扩展+广播实现逐行元素比较
result = tf.math.greater(X[:, :, tf.newaxis], X[:, tf.newaxis, :])

# 输出验证
print(result.numpy())
print("Result shape:", result.shape)

代码解释 #

• X[:, :, tf.newaxis]：将原n×m张量转换为n×m×1，相当于把每一行的元素都转为列向量形式
• X[:, tf.newaxis, :]：将原n×m张量转换为n×1×m，相当于把每一行保持行向量的同时扩展维度
• tf.math.greater：利用广播机制，自动对每个行的列向量与行向量进行逐元素比较，直接生成n×m×m的布尔张量，完全符合需求

输出结果 #

运行上述代码后，输出与示例完全一致：

[[[False False]
  [ True False]]

 [[False  True]
  [False False]]

 [[False False]
  [ True False]]]
Result shape: (3, 2, 2)

备选循环方案（不推荐） #

如果需要更直观的逐行处理逻辑，也可以用循环+堆叠的方式，但性能远不如向量化广播，仅适用于小张量场景：

matrices = []
for row in tf.unstack(X):
    # 对单行生成m×m比较矩阵
    mat = tf.math.greater(tf.reshape(row, (-1, 1)), row)
    matrices.append(mat)
# 沿第一个轴堆叠得到三维张量
result = tf.stack(matrices, axis=0)

内容的提问来源于stack exchange，提问作者Juan Castaño