loom_test.py 文件源码

def test_three_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()
      vals = [weaver(np.array([0, 1, 1 << k], dtype='int64'))
              for k in range(8)]
      for _ in xrange(3):
        vals = [weaver.add(*args) for args in group_values(vals, 2)]
      big_sum = vals[0]
      result = output_tensor.eval(
          feed_dict=weaver.build_feed_dict([big_sum]))
    self.assertTrue((result == np.array([[0, 8, 255]], dtype='int64')).all())