using while_loop over the tensor for creating a mask in tensorflow

Question

I want to create a mask with iterating over the tensor. I have this code:

import tensorflow as tf

out = tf.Variable(tf.zeros_like(alp, dtype=tf.int32))

rows_tf = tf.constant (
[[1, 2, 5],
 [1, 2, 5],
 [1, 2, 5],
 [1, 4, 6],
 [1, 4, 6],
 [2, 3, 6],
 [2, 3, 6],
 [2, 4, 7]])

columns_tf = tf.constant(
[[1],
 [2],
 [3],
 [2],
 [3],
 [2],
 [3],
 [2]])

I want to iterate through rows_tf and accordingly columns_tf to create a mask over the out.

for example, it will mask the index at [1,1] [2,1] and [5,1] in the out tensor equals 1.

for the second row in rows_tf indexes at [1,2] [2,2] [5,2] in the out tensor will be set to 1 and so on for the total 8 rows

So far I have done this, though it does not run successfully:

body = lambda k, i: (tf.add(out[rows_tf[i][k]][columns_tf[i][i]], 1)) # find the corresponding element in out tensor and add 1 to it (0+1=1)
k = 0
n2, m2 = rows_tf.shape
for i in tf.range(0,n2): # loop through rows in rows_tf    
    cond = lambda k, _: tf.less(k, m2) #this check to go over the columns in rows_tf
    tf.while_loop(cond, body, (k, i))

it raises this error:

TypeError: Cannot iterate over a scalar tensor. 
in this while cond(*loop_vars):

I have gone through several links namely here to make sure Im following the instruction, but could not fix this one.

Thanks for the help

Answer 1

You can do that without a loop using tf.scatter_nd like this:

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    out = tf.zeros([10, 4], dtype=tf.int32)
    rows_tf = tf.constant(
        [[1, 2, 5],
         [1, 2, 5],
         [1, 2, 5],
         [1, 4, 6],
         [1, 4, 6],
         [2, 3, 6],
         [2, 3, 6],
         [2, 4, 7]], dtype=tf.int32)
    columns_tf = tf.constant(
        [[1],
         [2],
         [3],
         [2],
         [3],
         [2],
         [3],
         [2]], dtype=tf.int32)
    # Broadcast columns
    columns_bc = tf.broadcast_to(columns_tf, tf.shape(rows_tf))
    # Scatter values to indices
    scatter_idx = tf.stack([rows_tf, columns_bc], axis=-1)
    mask = tf.scatter_nd(scatter_idx, tf.ones_like(rows_tf, dtype=tf.bool), tf.shape(out))
    print(sess.run(mask))

Output:

[[False False False False]
 [False  True  True  True]
 [False  True  True  True]
 [False False  True  True]
 [False False  True  True]
 [False  True  True  True]
 [False False  True  True]
 [False False  True False]
 [False False False False]
 [False False False False]]

Alternatively, you could also do this using boolean operations only:

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    out = tf.zeros([10, 4], dtype=tf.int32)
    rows_tf = tf.constant(
        [[1, 2, 5],
         [1, 2, 5],
         [1, 2, 5],
         [1, 4, 6],
         [1, 4, 6],
         [2, 3, 6],
         [2, 3, 6],
         [2, 4, 7]], dtype=tf.int32)
    columns_tf = tf.constant(
        [[1],
         [2],
         [3],
         [2],
         [3],
         [2],
         [3],
         [2]], dtype=tf.int32)
    # Compare indices
    row_eq = tf.equal(tf.range(out.shape[0])[:, tf.newaxis],
                      rows_tf[..., np.newaxis, np.newaxis])
    col_eq = tf.equal(tf.range(out.shape[1])[tf.newaxis, :],
                      columns_tf[..., np.newaxis, np.newaxis])
    # Aggregate
    mask = tf.reduce_any(row_eq & col_eq, axis=[0, 1])
    print(sess.run(mask))
    # Same as before

However this would in principle take more memory.