Операция трансляции по строкам 3D-тензора (Tensorflow)

Question

У меня есть тензор 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)

Answer 1

Попробуйте это:

result = tf.reduce_prod(self.A * p[:, tf.newaxis] + self.B, axis=2)