HIMANSHU RAI
HIMANSHU RAI

Reputation: 305

Is it possible to have multiple conditions defined in tf.while_loop

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

Answers (1)

nessuno
nessuno

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

Related Questions