我有以下代码
output = T.switch(cond, a, b)
其中cond是(N,1) bool Tensor，而a和b是(N, M)数字张量M非常大。条件以行方式运行。
我可以通过在T.repeat()上运行cond来轻松完成切换工作，但这很慢。有没有办法可以有效地使cond中的bool决定是否应该返回a或b？
答案 0 :(得分：3)
有没有办法可以有效地让cond中的bool决定是否应该返回a或b？
是的，你可以做到
cond * a + (1-cond) * b
cond将广播为(N, M)形状。
这应该接近理论极限，即内存带宽：此操作需要读取N*M个元素并写入N*M。
相反，我们阅读了2*N*M，但删除了条件逻辑。
（我不会让Theano出现在我面前，所以我不确定它是否比T.switch更快，但它应该和它一样好。而且，我和＃39; d尝试将cond投射到与dtype和a相同的b
如果您想就地更新a，可以使用T.set_subtensor执行此操作：
a = np.random.uniform(size=(N, M)).astype(np.float32)
b = np.random.uniform(size=(N, M)).astype(np.float32)
a = theano.shared(a)
b = theano.shared(b)
c = T.vector() # mostly 0, presumably (1-cond)
nz = T.nonzero(c)
s = T.set_subtensor(a[nz], b[nz])
fn = theano.function([c], [], updates=[(a, s)])
...
fn(1-cond)
它可能会或可能不会比第一种方法更快，具体取决于N，M和其他因素。