loom_test.py 文件源码

def test_two_layer_sum_network(self):
    shape = loom.TypeShape('int64', (3,))
    ops = {'add': BinaryLoomOp(shape, tf.add)}
    the_loom = loom.Loom(named_ops=ops)
    output_tensor = the_loom.output_tensor(shape)
    with self.test_session():
      weaver = the_loom.make_weaver()
      c1 = weaver(np.array([1, 2, 3], dtype='int64'))
      c2 = weaver(np.array([2, 4, 6], dtype='int64'))
      c3 = weaver(np.array([3, 6, 9], dtype='int64'))
      c4 = weaver(np.array([4, 8, 12], dtype='int64'))
      sum_1_2 = weaver.add(c1, c2)
      sum_3_4 = weaver.add(c3, c4)
      sum_1_2_3_4 = weaver.add(sum_1_2, sum_3_4)
      result = output_tensor.eval(
          feed_dict=weaver.build_feed_dict([sum_1_2_3_4]))
    self.assertTrue((result == np.array([[10, 20, 30]], dtype='int64')).all())