Adding consecutive elements across dimension in tensorflow (np.add.reduceat)

Question

Is there any tensorflow function or elegant of doing following task

consider I have 2d tensor with dimensions say $3 \times 10$.now I want to add consecutive elements along dimension $1$(or row) where number of consecutive elements to be added is decided by tensor.if it says $[2,2,4,2]$ the output tensor should be of size $ 3 \times 4$. because $ [ a_1 + a_2 , a_3+a_4 , a_5+a_6+a_7+a_8 , a_9+a_10]$ would tensor across each row

ex:

$\begin{bmatrix}
1 & 2 & 3 & 4 & 5 & 6 & 1 & 2 & 3 &4\\ 
7 & 8 & 9 & 10 &11  &12 & 7 & 8 & 9 & 10\\ 
13 &14  &15  &16  &17  &18 & 13 &14 &15 &16
\end{bmatrix}$

and output should be following

$\begin{bmatrix}
3 & 7 & 14 & 7\\ 
15 & 19 &38 &19\\ 
27  &31 & 62 &31
\end{bmatrix}$

EDIT:there seems to be function for this in numpy np.add.reduceat

Answer 1

There is no such function in tensorflow. But you can construct it by splitting the array:

def tf_reduceat(data, at_array, axis=-1):
    split_data = tf.split(data, at_array, axis=axis)
    return tf.stack([tf.reduce_sum(i, axis=axis) for i in split_data], axis=axis)


a = tf.constant([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4], 
                 [7, 8, 9, 10, 11, 12, 7, 8, 9, 10], 
                 [13, 14, 15, 16, 17, 18, 13, 14, 15, 16]])
result = tf_reduceat(a, [2, 2, 4, 2])

Running result yields:

array([[ 3,  7, 14,  7],                                                                                                       
       [15, 19, 38, 19],                                                                                                       
       [27, 31, 62, 31]], dtype=int32)