У меня есть тензор p размерности [BATCH_SIZE, 128] и тензоры A, B обоих измерений [528, 128]. Я хотел бы построить новый тензор p 'размера [BATCH_SIZE, 528], где столбец j определен как:
tf.reduce_prod(self.A[j,:] * p + self.B[j,:], axis=1)
В настоящее время у меня есть реализация этого перебора с использованием цикла for, но это очень медленно. Есть ли способ, которым я могу использовать вещание или что-то еще, чтобы ускорить это?
ps = []
for j in xrange(self.A.shape[0]):
a = self.A[j,:]
b = self.B[j,:]
ps.append(tf.reduce_prod(a * p + b, axis=1))
p = tf.stack(ps, axis=1)