Example ode_demo using y**3?

[pswpswpsw]

This example ode_demo.py under Example is considering a NN with input cubic nonlinearity. Since the original dynamics is y**3 so it seems a bit weird. This implies that as long as the NN is learning the linear matrix mapping, it will work.

class ODEFunc(tf.keras.Model):
    def __init__(self, **kwargs):
        super(ODEFunc, self).__init__(**kwargs)

        self.x = tf.keras.layers.Dense(50, activation='tanh',
                                   kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=0.1))
        self.y = tf.keras.layers.Dense(2,
                                   kernel_initializer=tf.keras.initializers.TruncatedNormal(stddev=0.1))

    def call(self, t, y):
        y = tf.cast(y, tf.float32)
        x = self.x(y ** 3)
        y = self.y(x)
        return y

Question

When I turned it back to y, so using NN to fully learning the dynamics, (see below) it doesn't converge and seems stuck at somewhere. Does anyone have success in using NODE for this simple example?