tensorflow2.1を用いた深層学習でのエラー(NotImplementedError)

前提・実現したいこと

tensorflow2.1を用いてNeural ODEモデルを構築し、入力が1次元、出力が1次元の回帰問題を解こうと思っております。

###該当のソースコード

Python
1import tensorflow as tf
2import numpy as np
3from scipy.integrate import solve_ivp
4from sklearn.utils import shuffle
5from sklearn.model_selection import train_test_split
6
7def get_dataset(N):
8    x = np.linspace(0., 1., N).reshape(-1, 1)
9    y = np.sin(np.pi * x)
10    return (x, y)
11
12def d_sigmoid(x):
13    return tf.multiply(tf.nn.sigmoid(x), tf.subtract(1., tf.nn.sigmoid(x)))
14
15class NeuralODEModel(tf.keras.Model):
16    def __init__(self, dim_in, dim_out, max_time, division, function, d_function):
17        super(NeuralODEModel, self).__init__(name='')
18        self.dim_in = dim_in
19        self.dim_out = dim_out
20        self.max_time = max_time
21        self.division = division
22        alpha = tf.Variable(tf.multiply(eps, tf.ones(shape=(self.division, self.dim_out), dtype=tf.dtypes.float32)))
23        beta = tf.Variable(tf.multiply(eps, tf.ones(shape=(self.division, self.dim_in, self.dim_in), dtype=tf.dtypes.float32)))
24        gamma = tf.Variable(tf.multiply(eps, tf.ones(shape=(self.division, self.dim_in), dtype=tf.dtypes.float32)))
25        self.params = (alpha, beta, gamma)
26        self.A = tf.constant(tf.eye(self.dim_out, self.dim_in, dtype=tf.dtypes.float32))
27        self.function = function
28        self.d_function = d_function
29        self.t = np.linspace(0., self.max_time, self.division)
30    
31    def call(self, x0):
32        def func(t, x, params, A, function, dim_in, dim_out, division):
33            x = x.astype(np.float32)
34            index = int(t * (division - 1))
35            if index > division - 1:
36                index = -1
37            x1 = x.reshape(-1, dim_in + dim_out)[:, :dim_in]
38            y1 = tf.add(tf.matmul(x1, tf.transpose(params[1][index])), params[2][index])
39            y2 = tf.multiply(params[0][index], function(tf.matmul(x1, tf.transpose(A))))
40            return tf.reshape(tf.concat([y1, y2], 1), [-1])
41        
42        y0 = np.zeros(shape=(len(x0), self.dim_out), dtype=np.float32)
43        solver = solve_ivp(func, (0., self.max_time), np.hstack([x0, y0]).reshape(-1), t_eval=self.t, args=(self.params, self.A, self.function, self.dim_in, self.dim_out, self.division))
44        self.x = solver.y.reshape(len(x0), self.dim_in+self.dim_out, -1)[:, :self.dim_in, :].astype(np.float32)
45        return solver.y.reshape(len(x0), self.dim_in+self.dim_out, -1)[:, self.dim_in:, -1]
46
47model = NeuralODEModel(1, 1, 1.0, 100, tf.nn.sigmoid, d_sigmoid)
48opt = tf.keras.optimizers.SGD()
49loss = tf.keras.losses.MeanSquaredError()
50train_loss = tf.keras.metrics.Mean()
51test_loss = tf.keras.metrics.Mean()
52n_data = 10000
53batch_size = 32
54epochs = 1000
55x, y = get_dataset(n_data)
56x_train, x_val, y_train, y_val = train_test_split(x, y, test_size=0.2)
57
58def gradient(model, x0, y_true):
59    def func(t, a, params, A, function, bT, x, division):
60        a = a.astype(np.float32)
61        index = int(t * (division - 1))
62        if index < 0:
63            index = 0
64        a1 = a.reshape(-1, model.dim_in)
65        a1 = tf.multiply(-1., tf.add(tf.matmul(a1, params[1][index]), tf.matmul(tf.multiply(bT, tf.multiply(params[0][index], function(tf.matmul(x[:, :, index], tf.transpose(A))))), A)))
66        return tf.reshape(a1, [-1])
67    
68    with tf.GradientTape() as tape:
69        y_pred = tf.Variable(model(x0))
70        error = loss(y_true, y_pred)
71    bT = tf.cast(tape.gradient(error, y_pred), tf.dtypes.float32)
72    aT = np.zeros_like(x0)
73    solver = solve_ivp(func, (model.max_time, 0.), aT.reshape(-1), t_eval=model.t[::-1], args=(model.params, model.A, model.d_function, bT, model.x, model.division))
74    a = solver.y.reshape(-1, model.dim_in, model.division)[:, :, ::-1]
75    g_alpha = tf.einsum('kj,kji->ij', bT, model.function(tf.einsum('jl,ilk->ijk', model.A, model.x)))
76    g_beta = tf.einsum('lji,lki->ijk', a, model.x)
77    g_gamma = tf.einsum('kji->ij', a)
78    return (g_alpha, g_beta, g_gamma)
79
80@tf.function
81def train_step(x, y):
82    opt.apply_gradients(zip(gradient(model, x, y), model.params))
83    error = loss(y, model(x))
84    train_loss(error)
85
86@tf.function
87def test_step(x, y):
88    y_pred = model(x)
89    error = loss(y, y_pred)
90    test_loss(error)
91
92for epoch in range(epochs):
93    x_train, y_train = shuffle(x_train, y_train)
94    for i in range(0, len(x_train), batch_size):
95        train_step(x_train[i:i+batch_size], y_train[i:i+batch_size])
96    for i in range(0, len(x_val), batch_size):
97        test_step(x_val[i:i+batch_size], y_val[i:i+batch_size])
98    
99    print("Epoch:{}  Train loss:{:.5f}  Test loss:{:.5f}".format(epoch+1, train_loss.result(), test_loss.result()))
100    train_loss.reset_states()
101    test_loss.reset_states()

公式ドキュメントを参考に構築してみましたが、以下のようなエラーが発生しております。
numpy配列とTensorが混在していることがエラーの原因だと思っているのですが、エラーの原因となる部分が分かりません。解決方法が分かる方いましたら、よろしくお願いいたします。

発生している問題・エラーメッセージ

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-549-a0b0dde6b671> in <module>()
      3     x_train, y_train = shuffle(x_train, y_train)
      4     for i in range(0, len(x_train), batch_size):
----> 5         train_step(x_train[i:i+batch_size], y_train[i:i+batch_size])
      6 

NotImplementedError: in converted code:

    <ipython-input-546-97616dd16c85>:3 train_step  *
        opt.apply_gradients(zip(gradient(model, x, y), model.params))
    <ipython-input-545-3763226feeb1>:14 gradient  *
        y_pred = tf.Variable(model(x0))
    c:\program files\python36\lib\site-packages\tensorflow_core\python\keras\engine\base_layer.py:778 __call__
        outputs = call_fn(cast_inputs, *args, **kwargs)
    <ipython-input-543-1184e5d2471c>:31 call  *
        solver = solve_ivp(func, (0., self.max_time), np.hstack([x0, y0]).reshape(-1), t_eval=self.t, args=(self.params, self.A, self.function, self.dim_in, self.dim_out, self.division))
    <__array_function__ internals>:6 hstack
        
    c:\program files\python36\lib\site-packages\numpy\core\shape_base.py:338 hstack
        arrs = atleast_1d(*tup)
    <__array_function__ internals>:6 atleast_1d
        
    c:\program files\python36\lib\site-packages\numpy\core\shape_base.py:67 atleast_1d
        ary = asanyarray(ary)
    c:\program files\python36\lib\site-packages\numpy\core\_asarray.py:138 asanyarray
        return array(a, dtype, copy=False, order=order, subok=True)
    c:\program files\python36\lib\site-packages\tensorflow_core\python\framework\ops.py:728 __array__
        " array.".format(self.name))

    NotImplementedError: Cannot convert a symbolic Tensor (neural_ode_model_23/Cast:0) to a numpy array.