Продукт Theano Dot

Question

Ниже приведен код для решения уравнения теплопроводности / диффузии с периодическими граничными условиями. Таким образом, при вычислении производной функции на каждом шаге РК мне нужно сделать скалярное произведение между входным вектором и матрицей Теплица. Это происходит правильно для первой итерации (проверено с результатами scipy.integrate), а для второй и далее что-то идет не так, и решение взрывается. Я не могу понять, в чем проблема. Пожалуйста помоги. Код является модификацией https://gist.github.com/michaelosthege/a75b565d3f653721fa235a07eb089912.

import abc
import theano.tensor as tt
import theano
import numpy as np
from scipy.linalg import inv, toeplitz

#################################### Integrator class #############################

class Integrator(object):
    __metaclass__ = abc.ABCMeta
    def step(self, t, dt, y, dydt, theta):
        raise NotImplementedError()

class TheanoIntegrationOps(object):
    def __init__(self, dydt_theano, integrator:Integrator):
        self.dydt_theano = dydt_theano
        self.integrator = integrator
        return super().__init__()

    def __step_theano(self, t, y_t, dt_t, nrk4steps, t_theta):
        return self.integrator.step(t, dt_t, y_t, self.dydt_theano, 
                    nrk4steps, t_theta)

    def __call__(self, y0, theta, times, nrk4steps):
        t_y0 = tt.as_tensor_variable(y0)
        t_theta = tt.as_tensor_variable(theta)
        t_times = tt.as_tensor_variable(times)
        dt = (times[1] - times[0])/nrk4steps

        Y_hat, updates = theano.scan(fn=self.__step_theano,
                                    outputs_info =[{'initial':t_y0}],
                                    sequences=[t_times],
                                    non_sequences=[dt, nrk4steps, t_theta],
                                    n_steps=len(times)-1)

        Y_hat = tt.concatenate((t_y0[None,:], Y_hat))
        Y_hat = tt.transpose(Y_hat)
        return Y_hat

# Simple RK integrator

class RungeKutta(Integrator):
    def step(self, t, dt, y, dydt, nrk4steps, theta):
        for i in range(nrk4steps.eval()):
            k1 = dt*dydt(y, t, theta)
            k2 = dt*dydt(y + 0.5*k1, t, theta)
            k3 = dt*dydt(y + 0.5*k2, t, theta)
            k4 = dt*dydt(y + k3, t, theta)
            y = y + (1./6.)*k1 + (1./3.)*k2 + (1./3.)*k3 + (1./6.)*k4
        return y


#################################################################

# The derivative function for heat equation
def dydt(c, t, theta):
    nx = 128
    gridx = np.arange(0.0, 1, 1.0/nx)
    dx = gridx[1] - gridx[0]
    z = np.zeros_like(gridx)
    z[0] = -2
    z[1] =  1 
    z[-1] = 1
    diff = 2
    diff_matrix = toeplitz(z) / dx**2
    diff_matrix = tt.as_tensor_variable(diff_matrix)
    dc_dt = tt.zeros_like(c)
    # This is what causes the problem
    dc_dt = theano.dot(c, diff_matrix)
    return dc_dt

if __name__=="__main__":
    # Number of RK4 iterations between each time point
    nrk4steps = 1
    # Time points
    times=np.linspace(0,1,101)
    # Initial Condition
    nx = 128
    gridx = np.arange(0.0, 1, 1.0/nx)
    T_y0 = np.sin(2*np.pi*gridx)
    # Diffusion constant
    T_theta = [2.0,]
    integrator = RungeKutta()
    Y_hat_TV = TheanoIntegrationOps(dydt, integrator)(
                T_y0, T_theta, times, nrk4steps)
    Ysim = np.transpose(Y_hat_TV.eval())
    print(Ysim)
~