Kiran K.
Kiran K.

Reputation: 366

Swapping elements within a matrix rows and columns - TensorFlow scatter_nd

I am trying to use the scatter_nd function in TensorFlow to reorder elements within rows of a Matrix. For example, suppose I have the code:

indices = tf.constant([[1],[0]])
updates = tf.constant([ [5, 6, 7, 8],
                        [1, 2, 3, 4] ])
shape = tf.constant([2, 4])
scatter1 = tf.scatter_nd(indices, updates, shape)
$ print(scatter1) = [[1,2,3,4]
                     [5,6,7,8]]

This reorders the rows of the updates matrix.

Instead of only being able to reorder the rows, I'd like to reorder the individual elements within each row as well. If I just have a vector (Tensor of rank 1), then this example works:

indices = tf.constant([[1],[0],[2],[3]])
updates = tf.constant([5, 6, 7, 8])
shape = tf.constant([4])
scatter2 = tf.scatter_nd(indices, updates, shape)
$ print(scatter2) = [6,5,7,8]

What I really care about is to be able to swap elements within each row in scatter1, as I had done in scatter2, but do it for each row of scatter1. I've tried various combinations of indices but keep getting errors that the sizes are inconsistent thrown by the scatter_nd function.

Upvotes: 1

Views: 3584

Answers (2)

Lerner Zhang
Lerner Zhang

Reputation: 7130

Suppose you want to swap elements in the second dimension either keeping the first dimension order or not.

import tensorflow as tf
sess = tf.InteractiveSession()


def prepare_fd(fd_indices, sd_dims):
    fd_indices = tf.expand_dims(fd_indices, 1)
    fd_indices = tf.tile(fd_indices, [1, sd_dims])
    return fd_indices

# define the updates
updates = tf.constant([[11, 12, 13, 14],
                       [21, 22, 23, 24],
                       [31, 32, 33, 34]])
sd_dims = tf.shape(updates)[1]

sd_indices = tf.constant([[1, 0, 2, 3], [0, 2, 1, 3], [0, 1, 3, 2]])
fd_indices_range = tf.range(0, limit=tf.shape(updates)[0])
fd_indices_custom = tf.constant([2, 0, 1])

# define the indices
indices1 = tf.stack((prepare_fd(fd_indices_range, sd_dims), sd_indices), axis=2)
indices2 = tf.stack((prepare_fd(fd_indices_custom, sd_dims), sd_indices), axis=2)

# define the shape
shape = tf.shape(updates)

scatter1 = tf.scatter_nd(indices1, updates, shape)
scatter2 = tf.scatter_nd(indices2, updates, shape)

print(scatter1.eval())

# array([[12, 11, 13, 14],
#        [21, 23, 22, 24],
#        [31, 32, 34, 33]], dtype=int32)

print(scatter2.eval())

# array([[21, 23, 22, 24],
#        [31, 32, 34, 33],
#        [12, 11, 13, 14]], dtype=int32)

May this example help.

Upvotes: 0

benj
benj

Reputation: 303

The following swaps the elements of each row of each row using scatter_nd

indices = tf.constant([[[0, 1], [0, 0], [0, 2], [0, 3]], 
                      [[1, 1], [1, 0], [1, 2], [1, 3]]])
updates = tf.constant([ [5, 6, 7, 8],
                        [1, 2, 3, 4] ])
shape = tf.constant([2, 4])
scatter1 = tf.scatter_nd(indices, updates, shape)
with tf.Session() as sess:
    print(sess.run(scatter1))

Giving an output of:
[[6 5 7 8] [2 1 3 4]]

The locations of the coordinate in indices define where the values are being taken from in updates and the actual cordinates define where the values will be placed in scatter1.

This answer is a few months late but hopefully still helpful.

Upvotes: 1

Related Questions