How do I sum a tensor by progressing index on specific axis?-CodePudding

I am working with time series returns and I want to find the cumulative return by time index. Assume the last index of a given time horizon should be itself, and the 1st index of a horizon should be the sum of all timesteps.

In the following example, I have a boolean mask indicating what timesteps should be added together for a given index. I just don't know where to go from here.

import tensorflow as tf

# Shape: (samples, horizon, feature)
source = tf.constant([[[0], [1], [2], [3]], [[4], [5], [6], [7]]])
# Shape: (horizon, horizon)
mask = tf.linalg.band_part(tf.ones((source.shape[1], source.shape[1])), 0, -1)

# Result should be tf.constant([[[6], [6], [5], [3]], [[22], [18], [13], [7]]])

CodePudding user response：

You can expand mask in the samples dimension, then use tf.matmul to get desired result

source = tf.constant([[[0], [1], [2], [3]], [[4], [5], [6], [7]]])
mask = tf.linalg.band_part(tf.ones((source.shape[1], source.shape[1])), 0, -1)
mask = tf.cast(mask, tf.int32)
ans = tf.matmul(mask[None,:,:], source)

Output

<tf.Tensor: shape=(2, 4, 1), dtype=int32, numpy=
array([[[ 6],
        [ 6],
        [ 5],
        [ 3]],

       [[22],
        [18],
        [13],
        [ 7]]])>

To do it without a mask, you can flip source on the horizon axis, perform a tf.cumsum, then flip the resultant tensor on the horizon axis again

ans = tf.reverse(tf.cumsum(tf.reverse(source, axis=[1]), axis=1), axis=[1])

Output

<tf.Tensor: shape=(2, 4, 1), dtype=int32, numpy=
array([[[ 6],
        [ 6],
        [ 5],
        [ 3]],

       [[22],
        [18],
        [13],
        [ 7]]])>