Tensorflow simple cumulative sum of product rnn cell

Hello i am trying to build a tensorflow model that calculates a cumulative sum of products of two of the input features, ie predicting on only (1,2) should return 2, and then predicting on (2,2) should give 6=(1 * 2) + (2 * 2)

model.predict([1,2])
>>> 2

model.predict([2,2])
>>> 6

model.reset_states()
model.predict([2,2])
>>> 4

i have tried the following:

import numpy as np
import tensorflow as tf


class MinimalRNNCell(tf.keras.layers.Layer):

    def __init__(self, units, **kwargs):
        self.states = np.array([0])
        self.state = np.array([0])
        self.units = units
        self.state_size = units
        super(MinimalRNNCell, self).__init__(**kwargs)

    def call(self, inputs, states):
        prev_output = states[0]
        output = tf.math.add(prev_output,inputs)
        
        return output, [output]
    

# Define model
#input
inp = tf.keras.layers.Input(shape=(2,))
#split input
x1,x2 = tf.split(inp, num_or_size_splits=2, axis=1)
#calculate product
product = tf.math.multiply(x1,x2)
#reshape product
time_product = tf.keras.layers.Reshape((1,1))(product)
#Define memory cell and layer
memory_product = MinimalRNNCell(units=1)
layer_product = tf.keras.layers.RNN(memory_product)
#calculate cumulative product
cumulative_product = layer_product(time_product)

output = cumulative_product

model = tf.keras.models.Model(inp, output)


if __name__=="__main__":
    x = np.array([
        [1, 2],
        [2, 2]
    ])

    model.compile()
    y = model.predict(x)
    print()
    print("outptut: ", y)
>>> [[2],
     [4]]

I feel like something like a cumulative sum is easily implementet using rnn or lstm cells but it does not work how i expect it to

source https://stackoverflow.com/questions/76348599/tensorflow-simple-cumulative-sum-of-product-rnn-cell