如何在NumPy/TensorFlow中按选择器提取Tensor元素？

当然可以！不管是NumPy还是TensorFlow，都有简洁的方法实现你描述的需求——根据二维选择器，从三维张量的对应位置提取指定索引的元素。先结合你的示例场景明确需求：

给定三维张量 T（形状为 N×N×K），以及二维选择器数组 selector（形状为 N×N），其中每个元素 selector[i,j] 指定提取 T[i,j] 中的第 selector[i,j] 个元素，最终得到一个 N×N 的结果数组。

先还原你的示例代码：

import numpy as np

# 示例三维张量
tensor = np.array([
    [[0.1, 0.8], [0.1, 0.8], [0.1, 0.8]],
    [[0.9, 0.3], [0.1, 0.8], [0.9, 0.3]],
    [[0.1, 0.8], [0.1, 0.8], [0.9, 0.3]]
])
# 二维选择器
selector = np.array([
    [0, 0, 1],
    [1, 1, 1],
    [1, 1, 0]
])
# 期望输出结果
want = np.array([
    [0.1, 0.1, 0.8],
    [0.3, 0.8, 0.3],
    [0.8, 0.8, 0.9]
])

NumPy 实现方案 #

你提到已经有Alok Singhal提供的解法，这里补充几种常用且直观的实现方式：

方法1：高级索引组合 #

利用NumPy的高级索引，将行、列索引与选择器索引组合，直接定位要提取的元素：

# 生成行索引（形状：N×1）和列索引（形状：1×N），通过广播匹配成N×N
rows = np.arange(tensor.shape[0])[:, None]
cols = np.arange(tensor.shape[1])

# 组合索引提取元素
result_np = tensor[rows, cols, selector]
# 验证结果是否匹配
print(np.allclose(result_np, want))  # 输出 True

方法2：使用np.take_along_axis（更简洁） #

这个函数专门用于沿指定轴提取对应位置的元素，只需要给选择器添加一个维度来对齐张量的最后一维：

# 将选择器从N×N扩展为N×N×1，与张量的K维度匹配
expanded_selector = selector[..., np.newaxis]
# 沿第2轴（K维度）提取元素，再去掉多余的维度
result_np = np.take_along_axis(tensor, expanded_selector, axis=2).squeeze(axis=2)

TensorFlow 实现方案 #

TensorFlow的实现思路和NumPy类似，同样可以通过索引组合或专用函数完成：

方法1：高级索引组合 #

import tensorflow as tf

# 转换为TensorFlow张量
tf_tensor = tf.convert_to_tensor(tensor)
tf_selector = tf.convert_to_tensor(selector)

# 生成行、列索引并广播匹配形状
rows = tf.range(tf_tensor.shape[0])[:, tf.newaxis]
cols = tf.range(tf_tensor.shape[1])

# 提取元素
result_tf = tf_tensor[rows, cols, tf_selector]
# 验证结果
print(tf.reduce_all(tf.abs(result_tf - want) < 1e-6))  # 输出 True

方法2：使用tf.gather_nd #

通过构建三维坐标数组，直接定位每个要提取的元素位置：

# 将行、列、选择器索引展平后组合成坐标数组（形状：N*N × 3）
indices = tf.stack([
    tf.reshape(rows, (-1,)),
    tf.reshape(cols, (-1,)),
    tf.reshape(tf_selector, (-1,))
], axis=1)
# 提取元素后还原为N×N形状
result_tf = tf.reshape(tf.gather_nd(tf_tensor, indices), selector.shape)

方法3：使用tf.gather配合维度扩展 #

类似NumPy的take_along_axis，用tf.gather沿指定轴提取：

# 扩展选择器维度到N×N×1
expanded_selector = tf.expand_dims(tf_selector, axis=-1)
# 沿第2轴提取元素，batch_dims指定对齐的前两维
result_tf = tf.gather(tf_tensor, expanded_selector, axis=2, batch_dims=2)
# 去掉多余维度
result_tf = tf.squeeze(result_tf, axis=-1)

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