Using While_loop Over The Tensor For Creating A Mask In Tensorflow

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 ( [[

Solution 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.