如何使用keras / tensorflow根据不同向量中的索引求和？

import tensorflow as tf
from keras import backend as K

input_tensor = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # 输入张量
index_tensor = tf.constant([[0, 1], [2, 0], [1, 2]])  # 索引张量

gathered_tensor = tf.gather(input_tensor, index_tensor)

sum_tensor = tf.reduce_sum(gathered_tensor, axis=1)

class IndexSumLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(IndexSumLayer, self).__init__()

    def call(self, inputs):
        input_tensor, index_tensor = inputs
        gathered_tensor = tf.gather(input_tensor, index_tensor)
        sum_tensor = tf.reduce_sum(gathered_tensor, axis=1)
        return sum_tensor

input_tensor = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # 输入张量
index_tensor = tf.constant([[0, 1], [2, 0], [1, 2]])  # 索引张量

output_tensor = IndexSumLayer()([input_tensor, index_tensor])