def test_single_fn(self):
with self.test_session() as sess:
n = tf.placeholder(tf.int64, [])
self.assertEqual(42, sess.run(plan._tf_nth([lambda: tf.constant(42)], n),
{n: 0}))
self.assertRaises(tf.errors.InvalidArgumentError, sess.run,
plan._tf_nth([lambda: tf.constant(42)], n), {n: 7})
def test_n_fns(self):
with self.test_session() as sess:
n = tf.placeholder(tf.int64, [])
nth = plan._tf_nth([lambda j=i: tf.constant(j) for i in xrange(4)], n)
for idx in xrange(4):
self.assertEqual(idx, sess.run(nth, {n: idx}))
self.assertRaises(tf.errors.InvalidArgumentError, sess.run, nth, {n: -1})