Reputation: 305
Is it possible to define to multiple conditions for termination of a tf.while_loop in tensorflow? For example depending on the two tensor values achieving two specific values. eg. i==2
and j==3
?
Also can I have several blocks of code in the body? In all the examples in the documentation, it seems that the body is more like a single statement returning a value or a tuple. I want to execute a set of several "sequential" statements in the body.
Upvotes: 1
Views: 6195
Reputation: 27042
tf.while_loop
accepts a generic callable (python functions defined with def
) or lambdas) that must return a boolean tensor.
You can, therefore, chain multiple conditions within the body of the condition using the logical operators, like tf.logical_and
, tf.logical_or
, ...
Even body
is a general python callable, thus you're not limited to lambdas and single statement functions.
Something like that is perfectly acceptable and works well:
import tensorflow as tf
import numpy as np
def body(x):
a = tf.random_uniform(shape=[2, 2], dtype=tf.int32, maxval=100)
b = tf.constant(np.array([[1, 2], [3, 4]]), dtype=tf.int32)
c = a + b
return tf.nn.relu(x + c)
def condition(x):
x = tf.Print(x, [x])
return tf.logical_or(tf.less(tf.reduce_sum(x), 1000), tf.equal(x[0, 0], 15))
x = tf.Variable(tf.constant(0, shape=[2, 2]))
result = tf.while_loop(condition, body, [x])
init = tf.global_variables_initializer()
with tf.Session() as sess:
sess.run(init)
print(sess.run(result))
Upvotes: 3