nairouz mrabah
nairouz mrabah

Reputation: 1217

Reducing a matrix by suming a fixed number of rows together in Tensorflow

I have a huge matrix with multiple rows. Actually, I want to add every 3 rows together to form a new matrix.

To better understand the problem, here is an example illustrating the desired output:

input  = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]
output = [[9, 12], [27, 30]]

I want to use tensorflow built-in operations in order to keep the graph differentiable.

Upvotes: 1

Views: 51

Answers (1)

benjaminplanche
benjaminplanche

Reputation: 15129

You could reshape your tensor in order to isolate the triplets in a new dimension, then tf.reduce_sum over that dimension:

import tensorflow as tf

x = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
shape_x = tf.shape(x)

# Reshaping the tensor to have triplets on dimension #1:
new_shape_for_reduce = tf.stack([shape_x[0] // 3, 3, shape_x[1]])
reshaped_x = tf.reshape(x, new_shape_for_reduce)
# Sum-reducing over dimension #1:
sum_3_rows = tf.reduce_sum(reshaped_x, axis=1)

with tf.Session() as sess:
    res = sess.run(sum_3_rows)
    print(res)
    # [[9  12]
    #  [27 30]]

Upvotes: 2

Related Questions