NeuralODE with multiple features? about deep-learning-notes HOT 3 CLOSED

import numpy as np
import numpy.random as npr
import tensorflow as tf
from tqdm import tqdm
import matplotlib.pyplot as plt
import tensorflow.contrib.eager as tfe

keras = tf.keras
tf.enable_eager_execution()
import neural_ode as node

class NNModule(tf.keras.Model):
    def __init__(self, num_filters):
        super(NNModule, self).__init__(name="Module")
        self.dense_1 = keras.layers.Dense(num_filters, activation="tanh")        
        self.t_dense_1 = keras.layers.Dense(num_filters, activation=lambda x: tf.sin(x))

    def call(self, inputs, **kwargs):
        # x = [batch_size, 3]
        # t = [1]
        t, x = inputs
                
        # t = [batch_size, 1]
        t = tf.tile(tf.reshape(t, [1, 1]), [12, 1])        
        # t1 = [batch_size, num_filters]
        t1 = self.t_dense_1(t)        
        # h = [batch_size, num_filters]
        h = self.dense_1(x) + t1
        return h

# create model and ode solver
model = NNModule(num_filters=3)
ode = node.NeuralODE(
    model, t=np.linspace(0, 1.0, 20),
    solver=neural_ode.rk4_step
)

x0 = tf.random_normal(shape=[12, 3])

with tf.GradientTape() as g:
    g.watch(x0)
    xN = ode.forward(x0)
    # some loss function here e.g. L2
    loss = xN ** 2

# reverse mode with adjoint method
dLoss = g.gradient(loss, xN)
x0_rec, dLdx0, dLdW = ode.backward(xN, dLoss)