如何实现兼容TensorFlow自动微分的双张量取值构建新张量？

第一步：向量化提取近似匹配元素（兼容自动微分） #

原循环实现会打断TensorFlow的自动微分链，下面用全向量化操作替代，核心是通过广播实现批量比较，再定位每个样本的首个匹配项：

# 1. 提取f中对应I索引的元素，得到形状为(n,)的张量f_I
f_I = tf.gather(f, I)

# 2. 扩展维度实现批量广播比较
f_I_expanded = tf.expand_dims(f_I, axis=1)  # 形状变为(n, 1)
D_flat = tf.reshape(D, (1, -1))             # 将(n,n)的D展平为(1, n²)

# 3. 批量计算近似匹配的布尔矩阵，形状为(n, n²)
matches = tf.experimental.numpy.isclose(f_I_expanded, D_flat, atol=1e-6)

# 4. 定位每个样本的首个匹配项：通过累积和找到第一个True的位置
cum_sum = tf.cumsum(tf.cast(matches, tf.float32), axis=1)
first_match_mask = tf.equal(cum_sum, 1.0)  # 标记每行第一个True的位置

# 5. 获取首个匹配项的索引并提取元素
first_match_indices = tf.argmax(tf.cast(first_match_mask, tf.float32), axis=1)
P_x = tf.gather(tf.reshape(D, (-1,)), first_match_indices)  # 形状为(n,)

关键说明： #

• 所有操作均为TensorFlow原生可微分操作，不会导致梯度断裂
• 用广播替代循环，大幅提升计算效率
• 依赖tf.cumsum定位首个匹配项，避免了循环中动态索引的问题

第二步：合并P_x与P_y为目标张量 #

直接使用tf.stack即可将两个形状为(n,)的张量按列合并，输出形状为(n,2)的张量，完全兼容自动微分：

# 假设P_y已通过tf.gather(g, indices)得到，形状为(n,)
P = tf.stack([P_x, P_y], axis=1)

效果验证： #

输出的P结构完全符合需求：

[[P_x[0], P_y[0]],
 [P_x[1], P_y[1]],
 ...,
 [P_x[n-1], P_y[n-1]]]

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