How to create a Keras layer from tf.math.segment

I would like to use the tf.math.segment_sum function in a Keras layer but I don't get the dimensions right.

As an example, I would like to sum the values of x_1 grouped by id in the dataframe df:

df = pd.DataFrame({'id':     [1, 1, 2, 2, 3, 3, 4, 4],
                   'x_1':    [1, 0, 0, 0, 0, 1, 1, 1],
                   'target': [1, 1, 0, 0, 1, 1, 2, 2]})

The 'model' I created looks as follows:

input_ = tf.keras.Input((1,), name='X')

cid = tf.keras.Input(shape=(1,), dtype='int64', name='id')

summed = tf.keras.layers.Lambda(lambda x: tf.math.segment_sum(x[0], x[1]), name='segment_sum')([input_, cid])

model = tf.keras.Model(inputs=[input_, cid], outputs=[summed])

I get an error about the rank:
ValueError: Shape must be rank 1 but is rank 2 for 'segment_sum/SegmentSum' (op: 'SegmentSum') with input shapes: [?,1], [?,1].

What do I do wrong here?

CodePudding user response：

I solved it using tf.gather. The working code is as follows:

input_ = tf.keras.Input((1,), name='X')

cid = tf.keras.Input(shape=(1,), dtype='int64', name='id')

summed = tf.keras.layers.Lambda(lambda x: tf.gather(tf.math.segment_sum(x[0], tf.reshape(x[1], (-1,))), x[1]), output_shape=(None,1), name='segment_sum')([input_, cid])

model = tf.keras.Model(inputs=[input_, cid], outputs=[summed])