I have two tensors as follows:
a = tf.constant([1, 0, 0, 1, 0, 1, 1])
b = tf.constant(0) # see below
I want to add one to b each time a[i] = 1 So I did this:
i = tf.constant(0)
def condition(i):
return i < 7
def f1(): return 1
def f2(): return 0
def body(i):
b = tf.constant(0)
b += tf.cond(tf.equal(a[i], 1), f1, f2)
return i + 1
r = tf.while_loop(condition, body, [i])
tf.print(r)
But I got the following error:
TypeError: Cannot iterate over a scalar-tensor.
I don't understand why after many tries to figure out and correct. Can someone help me fix this? Also, I wanted to retrieve the value of b at the end of the loop, how to do this?
Thank you in advance.
You can't have border effects in your loop. It means that your condition function and your body function must take as argument all the variable that you are using in your loop.
You should declare your functions the following way:
def condition(i,a,b):
return i < 7
def f1(): return 1
def f2(): return 0
def body(i,a,b):
b += tf.cond(tf.equal(a[i], 1), f1, f2)
return i + 1, a, b
Notice that the body function returns all the loop variables, in the same order that they were passed to the function.
and call the while loop passing the tuple (i,a,b) as your loop variables:
>>> tf.while_loop(condition, body, (i,a,b))
(<tf.Tensor: shape=(), dtype=int32, numpy=7>,
<tf.Tensor: shape=(7,), dtype=int32, numpy=array([1, 0, 0, 1, 0, 1, 1], dtype=int32)>,
<tf.Tensor: shape=(), dtype=int32, numpy=4>)
The loop returns a tuple containing the values of (i,a,b) after the execution of the loop. The last element is b, which is equal to 4, corresponding to the number of 1 in a.
Side note: I assume that your example is a simplified example, but if that's not the case, you should replace that while_loop with a call to tf.math.reduce_sum.