在两个不同批次的张量之间切换

Question

假设x和y是两个N-dimensional张量，两者都具有相同的尺寸，第一个尺寸的尺寸为S（批量大小）。假设b是一个1-dimensional张量的布尔值，大小为S。

我想生成z，一个N维张量定义为：

z[i] = b[i] ? x[i] : y[i] for i from 0 to (S-1)

其中x[i]指的是x的第i个(N-1)-dimensional切片。

最简单的方法是什么？我认为tf.cond可行，但它只接受标量值谓词。谢谢！

Answer 1

tf.where应该有效，并且支持这种广播。如果您发现自己想要一个条件执行的批处理版本（其中一个或两个分支的计算成本很高），那就是also possible。