Test if RNN can learn the length of input sequence

TestRNNSequenceCount.py

import tensorflow as tf
import random


def genSequence():
    len = 101
    X = []
    for i in range(1, len):
        X.append([[random.random()]]*i)
    y = []
    for i in range(1, len):
        y.append((i * 1.0)/len)
    return (X, y)


if __name__ == "__main__":

    X, y = genSequence()

    XTensor =  tf.ragged.constant(X, dtype=tf.float32)
    yTensor = tf.ragged.constant(y, dtype=tf.float32)

    print(XTensor)
    print(yTensor)

    # Build the Keras model.
    keras_model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(None, 1), dtype=tf.float32, ragged=True),
        tf.keras.layers.LSTM(32, use_bias=True),
        # MultiHeadAttention(head_num=1, name='Multi-Head'),
        # tf.keras.layers.Dense(32, activation='sigmoid'),
        tf.keras.layers.Dense(1)
    ])
    keras_model.summary()

    # keras_model.compile(loss='binary_crossentropy', optimizer='adam')
    keras_model.compile(loss='MSE', optimizer=tf.keras.optimizers.Adam(0.001))
    keras_model.fit(XTensor, yTensor, epochs=200)
    print(keras_model.predict(tf.constant([[[1]]])))
    print(keras_model.predict(tf.constant([[[1]]*50])))
    print(keras_model.predict(tf.constant([[[1]]*150])))