CODEWITHSUNDEEP
pythontensorflow
Tom
Tom

Reputation: 357

Summing columns of a 2D tensor according to multiple sets of indices

In tensorflow, I would like to sum columns of a 2D tensor according to multiple sets of indices.

For example:

Summing the columns of the following tensor

[[1 2 3 4 5]
 [5 4 3 2 1]]

according to the 2 sets of indices (first set to sum columns 0 1 2 and second set to sum columns 3 4)

[[0,1,2],[3,4]]

should give 2 columns

[[6  9]
 [12 3]]

Remarks:

  1. All columns' indices will appear in one and only one set of indices.
  2. This has to be done in Tensorflow, so that gradient can flow through this operation.

Do you have any idea how to perform that operation? I suspect I need to use tf.slice and probably tf.while_loop.

Upvotes: 4

Views: 940

Answers (2)

javidcf
javidcf

Reputation: 59701

You can do that with tf.segment_sum:

import tensorflow as tf

nums = [[1, 2, 3, 4, 5],
        [5, 4, 3, 2, 1]]
column_idx = [[0, 1, 2], [3, 4]]

with tf.Session() as sess:
    # Data as TF tensor
    data = tf.constant(nums)
    # Make segment ids
    segments = tf.concat([tf.tile([i], [len(lst)]) for i, lst in enumerate(column_idx)], axis=0)
    # Select columns
    data_cols = tf.gather(tf.transpose(data), tf.concat(column_idx, axis=0))
    col_sum = tf.transpose(tf.segment_sum(data_cols, segments))
    print(sess.run(col_sum))

Output:

[[ 6  9]
 [12  3]]

Upvotes: 2

troymyname00
troymyname00

Reputation: 702

I know of a crude way of solving this in NumPy if you don't mind solving this problem with NumPy.

import numpy as np

mat = np.array([[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]])

grid1 = np.ix_([0], [0, 1, 2])
item1 = np.sum(mat[grid1])
grid2 = np.ix_([1], [0, 1, 2])
item2 = np.sum(mat[grid2])
grid3 = np.ix_([0], [3, 4])
item3 = np.sum(mat[grid3])
grid4 = np.ix_([1], [3, 4])
item4 = np.sum(mat[grid4])

result = np.array([[item1, item3], [item2, item4]])

Upvotes: -1

Related Questions