TensorFlow中如何实现类似Torch select的子张量提取？有无更优方法？

嘿，这个问题我太有共鸣了！刚从Torch转TensorFlow的时候，也纠结过怎么用更简洁的方式替代select。其实TensorFlow里有几种比tf.gather更直观的方法，完全不用“大材小用”：

1. 直接索引语法（最推荐） #

TensorFlow支持和NumPy一致的索引方式，这是最简洁的实现，完美对应Torch的select逻辑：

• 对应Torch的tensorA:select(0, 0)（提取第0个维度的第0个元素）：

  import tensorflow as tf
  
  tensorA = tf.constant([[0,1,0,1],[1,0,1,0]])
  # 直接索引第0行，省略后面的冒号也可以
  result = tensorA[0]  # 或者 tensorA[0, :]
  print(result.numpy())  # 输出: [0 1 0 1]
  

• 对应Torch的tensorA:select(1, 1)（提取第1个维度的第1个元素）：

  result = tensorA[:, 1]
  print(result.numpy())  # 输出: [1 0]
  

这种方式不仅代码更短，可读性也极强，完全不需要额外的API调用，特别适合这种单索引提取的场景。

2. 保持维度的场景 #

如果需要和tf.gather(indices=[i], axis=d)一样保持原维度（比如返回形状为(1,4)而不是(4,)的张量），可以用切片语法：

# 对应select(0,0)且保持维度
result = tensorA[0:1, :]
print(result.shape)  # 输出: (1, 4)

# 对应select(1,1)且保持维度
result = tensorA[:, 1:2]
print(result.shape)  # 输出: (2, 1)

3. 关于tf.gather的补充 #

你提到的tf.gather确实是通用的高维索引工具，适合批量提取多个索引的场景，但对于单个索引的提取，直接索引语法显然更高效简洁。不过如果是动态维度或者需要在函数中通过变量指定维度和索引，tf.gather依然是可靠的选择——比如：

d = 0
i = 0
result = tf.gather(tensorA, indices=i, axis=d)  # 这里indices可以是单个值，不用传列表

注意这里indices可以直接传单个整数，不需要用列表包裹，代码也能简化不少。

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