我有两个变量A和B，其形状分别为（N，M）和（M，K）。
我可以计算C = tf.matmul(A, B)并得到形状为（N，K）的新矩阵C。
或者，我可以将其重写为C = tf.reduce_sum(tf.expand_dims(A, -1) * tf.expand_dims(B, 0), 1)。显然，该表达式速度较慢，需要更多的内存来计算。
我想计算C = tf.reduce_max(tf.expand_dims(A, -1) * tf.expand_dims(B, 0), 1)。 Kinda将两个矩阵相乘，但设置新元素C i，j = max k （A i，k * B k，j ）。有没有类似于tf.matmul的优化形式？