loom_test.py 文件源码

def test_two_ops_network(self):
    shape = loom.TypeShape('int64', (3,))
    ops = {'add': BinaryLoomOp(shape, tf.add),
           'mul': BinaryLoomOp(shape, tf.multiply)}
    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'))
      sum_2_3 = weaver.add(c2, c3)
      sum_12_13 = weaver.mul(c1, sum_2_3)
      result = output_tensor.eval(
          feed_dict=weaver.build_feed_dict([sum_12_13]))
    self.assertTrue((result == np.array([[5, 20, 45]], dtype='int64')).all())