loom_test.py 文件源码

def test_gradient(self):
    x_var = tf.Variable(tf.zeros([3], dtype='float64'), name='x')
    shape = loom.TypeShape('float64', (3,))
    ops = {'add': BinaryLoomOp(shape, tf.add),
           'mul': BinaryLoomOp(shape, tf.multiply)}
    the_loom = loom.Loom(named_tensors={'x': x_var}, named_ops=ops)

    output_tensor = the_loom.output_tensor(shape)
    output = tf.reduce_sum(output_tensor)
    gradient = tf.gradients(output, [x_var])[0]
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())

      weaver = the_loom.make_weaver()
      m = weaver(np.array([1, 2, 3], dtype='float64'))
      b = weaver(np.array([47, 9, -1], dtype='float64'))
      mx = weaver.mul(m, weaver.x)
      mx_plus_b = weaver.add(mx, b)
      result = gradient.eval(feed_dict=weaver.build_feed_dict([mx_plus_b]))
    self.assertTrue((result == np.array(
        [1.0, 2.0, 3.0], dtype='float64')).all())