tensorflow中tensor,从每行取指定索引元素
tensorflow
一个面向所有人的开源机器学习框架
项目地址:https://gitcode.com/gh_mirrors/te/tensorflow

·
实验有需求,需要对tensor中每一行取一个不同的索引的元素,其中tensor为2维(本文方法适合任意维),因此本文以2维tensor为例。
# 二维tensor
g = tf.constant([[1,2,3,4,5,6,7,8],[9,8,7,6,5,4,3,2]])
# 每一行取的index,在本例中,正确取值为[3, 2],即第一行index=2的元素和第二行index=7的元素
h_index = np.array([2, 7]).reshape(-1, 1)
# 构建一个numpy的arange列表,其长度为tensor的行数
line = np.arange(2).reshape(-1, 1)
# 注意上面两个numpy数组的格式都是(-1, 1)
# 将h_index和line合并
index = np.hstack((line, h_index))
# 使用tf.gather_nd来取值
result = tf.gather_nd(g, index)
如上即可,返回仍为tensor




一个面向所有人的开源机器学习框架
最近提交(Master分支:27 天前 )
4f64a3d5
Instead, check for this case in `ResolveUsers` and `ResolveOperand`, by querying whether the `fused_expression_root` is part of the `HloFusionAdaptor`.
This prevents us from stepping into nested fusions.
PiperOrigin-RevId: 724311958
27 天前
aa7e952e
Fix a bug in handling negative strides, and add a test case that exposes it.
We can have negative strides that are not just -1, e.g. with a combining
reshape.
PiperOrigin-RevId: 724293790
27 天前
更多推荐
所有评论(0)