import tensorflow as tf

    mask3 = tf.constant([[1, 1],
                         [1, 1]], dtype=tf.float32)
    mask4 = tf.constant([[0, 0],
                         [0, 0]], dtype=tf.float32)
    mask5 = tf.constant([[1, 0],
                         [1, 0]], dtype=tf.float32)
    masks = tf.stack([mask0, mask1, mask2, mask3, mask4, mask5])
    scores = tf.constant([[0.05, 1.0, 0.2],
