如何使用tensorflow创建产品图层？

Question

我正在尝试使用tensorflow实现SPN（http://spn.cs.washington.edu/pubs.shtml）。 在SPN中，我们只有sum或product节点。通过对其相应边缘权重和对应子项的乘积求和来计算和节点的值。产品节点的值是其子节点的乘积。

这是一个玩具示例：

我可以使用这段张量流代码轻松计算第一个总和层：

import tensorflow as tf
import numpy as np
tf.set_random_seed(0)

X = tf.placeholder(tf.float32, [None, 6])
W = tf.Variable(tf.ones([6, 4]))
S = tf.matmul(X,W)


sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)
r = sess.run(S, feed_dict={X: np.ones([1,6])})
print(r)

这为我提供了第一层中和节点的正确值：

[[ 6.  6.  6.  6.]]

我的问题是如何计算产品层？

产品节点应创建以下输出：

[[36. 36.]]

我尝试使用tf.reduce_prod，但它只提供了[[1296.0]]。

感谢您的想法。