def test_pfor_with_closure_multi_out(self):
val = np.arange(7.)[:, np.newaxis]
tf_val = tf.constant(val)
def tf_fn(x):
return tf.gather(tf_val, x)**2, tf.gather(tf_val, x)
def np_fn(x):
return nptf.gather(val, x)**2, nptf.gather(val, x)
self.assertAllEqual(
self.evaluate(tf_pfor.pfor(tf_fn, 7)),
np_pfor.pfor(np_fn, 7))
def test_pfor(self):
self.assertAllEqual(
self.evaluate(tf_pfor.pfor(lambda x: tf.ones([]), 7)),
np_pfor.pfor(lambda x: nptf.ones([]), 7))