<jupyter_start><jupyter_text>Tensorflow Experiments on Neural Ordinary Differential Equations> You can contact me on twitter as [@mandubian](http://twitter.com/mandubian)This notebook is a sandbox to test concepts exposed in this amazing paper:> Neural Ordinary Differential Equations> http://arxiv.org/abs/1806.07366>> Authors: >    Chen, <NAME>.>    <NAME>>    <NAME>>    Duvenaud, DavidMy idea is to reproduce the concepts exposed in the paper fully in Tensorflow using Eager-Mode and GradientTape. I didn't want to depend on Autograd and to be able to use classic Keras models.For ODE Solver, I wanted to compare implementations in TF with Scipy very robust ones (the only one I found is Runge-Kutta Dopri5). A few preliminary ideas (to let you scratch your head with me)> Skip it if you just want to see it running Modeling the derivative of hidden stateDeep neural networks as we know in general are aimed at simulating a function $y = f(x)$ from a set of known inputs/outputs (a supervised learning based on a groundtruth). The classic approach is to represent the function with inputs, outputs and a _hidden state_ we don't know but want to learn. This hidden state is represented by a sequence of discrete deep neural layers whose internal parameters are learnt by step by step back-propagation of gradients on a chosen derivable loss function and an optimization mechanism to reach the best minimum (ideally global). So the hidden state to be learnt is parameterized by the trainable parameters of those discrete layers.>In _Neural ODE_ paper (http://arxiv.org/abs/1806.07366), the authors propose a new approach by modeling not the hidden state of the function but the derivative of the hidden state. Residual layer as Ordinary Differential EquationThe authors make this assumption starting from the case of well-known residual neural networks _ResNet_ constituted of multiple residual layers. A residual layer is meant to _keep in memory_ what you knew in previous layer and to slighly update it without forgetting it. A residual layer can be represented by the following equation:$$z_{t+1} = z_t + f(z_t, {\theta}_t) \\$$where:- $z_t$ is the input of the layer- $z_{t+1}$ is the output of the layer- $f$ is the function parameterized by $θ_t$ that we want to learn- $t$ can be seen as the depth of this layer in a full ResNet with $0, 1, ..., t, t+1, ..., t+n$ successive residual layers.Previous equation can be rewritten as:$$z_{t+1} - z_t = f(z_t, {\theta}_t) * ({\delta}_t = 1) \\or \\\frac {z_{t+1} - z_t} {({\delta}_t = 1)} = f(z_t, {\theta}_t)$$${\delta}_t = 1$ can be seen as discrete depth increment going from one residual layer to the next in the ResNet.Let's now consider this depth $t$ is no more discrete but continuous. We can easily infer the following:$$\begin{aligned}{\frac {dz(t)}{dt}} = f(z(t), t, {\theta})\end{aligned}$$The previous equation is called an _ordinary differential equation (ODE)_. There is a whole field of study in mathematics about solving ordinary differential equations, with so-called _ODE Solvers_, i.e solving equations of kind:$$\text{Find } z(t) \text{ knowing }\begin{aligned}{\frac {dz(t)}{dt}} & = f(z(t), t) \\\end{aligned}\text{ given an initial value }\begin{aligned}z(t_{0}) & = z_{0} \\\end{aligned} \text{ and a discretized interval }\begin{aligned}{t_0, t_1, ... t_n}\end{aligned}$$ODE Solvers are numerical methods to compute fast but robust approximated integrals (analytical methods exist but can't solve many cases of ODE like $\dot{y} = sin(x^2 + y^2)$). During your math study, you might have seen this Euler method https://en.wikipedia.org/wiki/Euler_method of approximating a function by using a naive construction with the tangent, an incremental step and an initial point on the curve. I'm not a mathematician so I won't teach you anything about ODE Solvers, you just need to know that this field has been studied and used intensively in multiple real-life apps (physics, geology, meteo, finance, chemistry etc...) and provide very efficient, optimized algorithms implemented in multiple libs and for example Scipy provides nice implementations of different ODE solvers. Computing hidden state using ODE SolverLet's imagine we have a ResNet with multiple residual layers of depths $0, 1, ..., t, t+1, ..., t+n$. In a classic ResNet, all layers have their own function $f_t$ but if we mutualize this function for all layers, we can generalize to the whole network:$$\begin{aligned}{\frac {dz(t)}{dt}} & = f(z(t), t)\hspace{0.5cm}\forall{t} \in [t_0, t_1, ...t_n]\end{aligned}$$So the $f$ function of this ResNet can be seen as the derivative of the hidden state ${\frac {dz(t)}{dt}}$ we are trying to learn.Using ODE Solver, if we know $f$ and an initial value $z(t_{0})$ and this interval discretization ${t_0, t_1, ..., t_n}$, we can compute a robust approximation $z(t)$ in the following way:$$\begin{aligned}z(t_n) & = ODESolve(z(t_0), f, [t_0, t_1, ...t_n], θ) \end{aligned}$$where:- $t_0$ is the initial layer- $[t_0, t_1, ...t_n]$ is the discretized interval- $z(t_0)$ is the initial value of the function- $z(t_n)$ is the output value of our network- $t_n$ is the output layer- ${\theta}$ are the parameters of the function $f$ we want to learn Depth as an input parameter of neural model$[t_0, t_1, ...t_n]$ represents the depth distribution of the layers of our _mono-function_ ResNet. In a classic network, the depth in terms of layers is a static part of the architecture of the network which can't be changed without retraining completely the network.In our case, depth distribution becomes a continuous space that we discretized and provide as a dynamic parameter to the ODE Solver at training or inference. We can imagine having different depth distribution at training or inference to speed computations for example. We could even imagine using more complicated depth distribution like non linear distribution or even learning an optimized shape for a given problem. Next studies...There is also an immediate analogy between this depth distribution and time distribution. It also quickly reminds of time-series issues and recurrent neural networks... They speak about it in the paper but I still haven't implemented that part... Next steps... Back-propagation by backward ODE SolvingSo that's cool but how do you represent and learn those ${\theta}$???I won't go in deep demonstration as it would be too long. Don't worry if you don't catch this part :DIt starts with a classic neural network representing our derivative $f$ with internal parameters $\theta$.Let simplify to a 1-step time interval $[t_0, t_1]$, an initial input value $z(t_0)$ and our known output $groundtruth(t_1)$.First, it computes current estimated value $z({t_1})$ by a classic forward computation with a call to our ODE Solver using current $\theta$ and input value $z(t_0)$.$$\begin{aligned}z(t_1) = ODESolve(z(t_0), f, [t_0, t_1], {\theta})\end{aligned}$$Then, it computes the classic Loss function:$$\begin{aligned}L(z(t_1), groundtruth(t_1)) & = L(ODESolve(z(t_0), f, [t_0, t_1], {\theta}), groundtruth(t_1))\end{aligned}$$The authors then consider the so-called _adjoint state_ (adjoint state method is the math _trick_ at the origin of backward auto-differentiation _aka_ gradient back-propagation that often simplifies gradient computation as neural networks are often many-to-scalars functions with lower gradient complexity in backward mode... This is another topic)$$\begin{aligned}a(t) = − {\frac {\partial L(z(t))}{\partial z(t)}}\end{aligned}$$The paper appendix proves that:$$\begin{aligned}{\frac {da(t)}{dt}} = - a(t){\frac {\partial{f(z(t), t, \theta)}}{\partial z}}\end{aligned}$$and the same respectively for ${\frac {da(t)}{d\theta}}$ and ${\frac {da(t)}{dt}}$${\frac {\partial{f(z(t), t, \theta)}}{\partial z}}$ is the classic gradient we already compute in normal residual layer using back-propagation.Knowing all the derivatives of $a(t)$ with respect to $z(t), t, \theta$, we can then compute $a(t)$ with a call to ODE Solver going backward in time from $t_1$ to $t_0$. The initial value is the gradient of loss ${\frac {\partial L(z(t_1))}{\partial z}}$ at $t_1$ between our current output $z(t_1)$ and our groundtruth...$$\begin{aligned}s0 & = [z(t1), {\frac {\partial L}{\partial z(t_1)}}, 0, −{\frac {\partial L}{t_1}}] \\[z(t_0), {\frac {\partial L}{\partial z(t_0)}}, {\frac {\partial L}{\partial θ}}, {\frac {\partial L}{\partial t}}] & = ODESolve(z(t_0), augmented_f, [t_1, t_0], {\theta})\end{aligned}$$By using a bit more machinery (an _augmented state_ grouping $z(t)$, ${\frac {da(t)}{dz}}$, ${\frac {da(t)}{d\theta}}$, ${\frac {da(t)}{dt}}$ together), the paper demonstrates how all required backward gradients can be computed in one single call to ODESolver and then you can train your parameters using a classic optimizer like SGD or Adam...Not so clear? It's normal :D ConclusionBy representing the residual-like hidden-state through its derivative and a _continuous_ depth distribution, it's possible to train model parameters with ODE Solvers. The hidden-state of the derivative is still represented by a static human-engineered discrete sequence of neural layers but the depth distribution becomes an input parameter of the network that can be changed dynamically at training or inference.Is it going to change your life as a datascientist drastically? Not yet... But it's bringing new cool ideas on how we can represent neural networks in another way, how it learns, how residual layers work. It opens new fields of study, drawing bridges between Deep Learning and existing older branches of mathematics. It also bring diversity in DL domain that IMHO tends to re-use the same tricks without always catching how they work. Other sources of inspirationThere are more precise details in this nice article https://rkevingibson.github.io/blog/neural-networks-as-ordinary-differential-equations/ and then read the paper first pages, it's quite simple after all. If I can understand it, you can too...This notebook is also very inspired by this cool article on Neural ODE & appealing Julia:https://julialang.org/blog/2019/01/fluxdiffeqIf you want to know more about ODE Solvers, there is a whole litterature about it with lots of different methods like `Runge-Kutta` or `LSODA` used in this notebook and you can find more details in scipyhttps://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.htmlscipy.integrate.odeint  Code under MIT License```Copyright 2019 <NAME>Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.```  Initialization<jupyter_code>from __future__ import absolute_import
from __future__ import print_function

import time
import os
import numpy as np
import scipy.integrate
from plotly import __version__
from plotly.offline import download_plotlyjs, init_notebook_mode, plot, iplot
#import plotly.plotly as py
import plotly.graph_objs as go
import plotly.io as pio

import tensorflow as tf
print("tensorflow version: ", tf.__version__)

tf.enable_eager_execution()
init_notebook_mode(connected=True)
config = tf.ConfigProto()
config.gpu_options.allow_growth=True

os.environ['CUDA_VISIBLE_DEVICES'] = '0'<jupyter_output><empty_output><jupyter_text>Test ODE Solvers (TF and Scipy) Let a derivative function:$\hspace{1cm}f'(t) = 2t$Let a set of values $t \in [0, 4.0]$Let the initial value be $f(t_0) = 0.0$A good ODE Solver should integrate to a parabolic:$\hspace{1cm}f(t) = t^2$  Test Full-TF ODE Solver (Runge-Kutta - Adaptative - Explicit - Dopri5)> `Dopri5` is an explicit runge-kutta method of order (4)5 due to Dormand & Prince (with stepsize control and dense output).>> Authors:> <NAME> and <NAME> Universite de Geneve, Dept. de Mathematiques CH-1211 Geneve 24>> **Please note that this implementation works well but the method tends to be less stable computationally speaking than the scipy LSODA implicit method used below**<jupyter_code>ts = np.linspace(0., 4.0, 10)

ys, info = tf.contrib.integrate.odeint(lambda y, t: tf.to_float(2*t), 0.0, ts, rtol = 0.01, full_output=True)

iplot([go.Scatter(x=ts, y=ys.numpy())])<jupyter_output><empty_output><jupyter_text>Test Scipy ODE Solver (LSODE)> By default Scipy is using on LSODA Solver method (there are other methods, even DOPRI5)>> Real-valued Variable-coefficient Ordinary Differential Equation solver, with fixed-leading-coefficient implementation. It provides automatic method switching between implicit Adams method (for non-stiff problems) and a method based on backward differentiation formulas (BDF) (for stiff problems).>> Source: http://www.netlib.org/odepack>> **Please note that this implementation is more stable computionally speaking than full-TF implementation above**<jupyter_code>ts = np.linspace(0., 4.0, 10)

ys = scipy.integrate.odeint(lambda y, t: 2*t, [0.0], ts)

iplot([go.Scatter(x=ts, y=np.squeeze(ys))])<jupyter_output><empty_output><jupyter_text>Illustrating ODE Solvers with oscillators  Two-dimensional damped oscillatorhttp://www.dam.brown.edu/people/mraissi/research/7_multistep_neural_networks/Let us consider the two-dimensional damped harmonic oscillator with cubic dynamics:$\begin{align}{\frac {dx} {dt}} & = −0.1x^3 + 2.0y^3 \\{\frac {dy} {dt}} & = −2.0x^3 − 0.1y^3 \\z & = (x, y)\end{align}$<jupyter_code># z = [x, y]
def damped_oscillator_2d(z, t, A):
    return np.dot(z**3, A)

A = np.array([
    # −0.1x^3 + 2.0y^3
    [-0.1, 2.0],
    # −2.0x^3 − 0.1y^3
    [-2.0, -0.1]])

# Initial u0 = [x0, y0]
z0 = np.array([2., 0.]).T

# Time sampling
N = 100  # Dataset size
max_T = 6.0 # Time sampling
ts = np.linspace(0., max_T, N)<jupyter_output><empty_output><jupyter_text>TF-RungeKutta ODE Solver<jupyter_code>zs, info = tf.contrib.integrate.odeint(lambda z0, ts: tf.tensordot(z0**3, A, axes=1), z0, ts,
                                method="dopri5", options={"max_num_steps":10, "first_step": 0.1}, rtol = 0.001,
                                full_output=True)
zs = zs.numpy()
xs = zs[:, 0]
ys = zs[:, 1]
iplot([go.Scatter(x=ts, y=xs, name="x"), go.Scatter(x=ts, y=ys, name="y")])
iplot([go.Scatter(x=xs, y=ys)])<jupyter_output><empty_output><jupyter_text>Scipy ODE Solver<jupyter_code>zs = scipy.integrate.odeint(damped_oscillator_2d, z0, ts, args=(A,))

xs = zs[:, 0]
ys = zs[:, 1]
iplot([go.Scatter(x=ts, y=xs, name="x"), go.Scatter(x=ts, y=ys, name="y")])
iplot([go.Scatter(x=xs, y=ys)])<jupyter_output><empty_output><jupyter_text>Helpers to call ODE Solvers from TF<jupyter_code>def ode_tf(f, z0, ts, args):
    """calls Full TF ODE Runge-Kutta Dopri5 Solver
    
    Args:
        f: a function of type (z0, ts, args) -> ys
            where z0, the initial value
                  ts, the time values
                  args, the supplementary args

        z0: initial value
        ts: time values on which to solve equation
        args: supplementary arsg
        
    Returns:
        array of solutions computed by ODE Solver
    Please note, TF solver params are hard-coded here but they could be managed in a cleaner way
    """
    # current TF odeint implementation doesn't manage args so we need this lambda
    zs, info = tf.contrib.integrate.odeint(
        lambda z0, ts: f(z0, ts, args), z0, ts,
        method="dopri5",
        options={"max_num_steps":1000, "first_step": 0.001},
        rtol = 0.001,
        full_output=True
    )
    return zs

def ode_scipy(f, z0, ts, args):
    """calls Python Scipy ODE Solver from TF
    
    Args:
        f: a function of type (z0, ts, args) -> ys
            where z0, the initial value
                  ts, the time values
                  args, the supplementary args

        y-z0: initial value
        ts: time values on which to solve equation
        args: supplementary args

    Returns:
        array of solutions computed by ODE Solver
    """
    # managing args as None... this is ugly but simplifies a bit later code
    if args is not None:
        r = tf.py_func(
            lambda z0, ts, args: scipy.integrate.odeint(f, z0, ts, args=(args,)),
            [z0, ts, args], z0.dtype)
    else:
        r = tf.py_func(
            lambda z0, ts: scipy.integrate.odeint(f, z0, ts, args=(None,)),
            [z0, ts], z0.dtype)
    return r<jupyter_output><empty_output><jupyter_text>Lotka–Volterra equations (the funny predators/preys or rabbits/foxes equations)> This sample was taken from article https://julialang.org/blog/2019/01/fluxdiffeqhttps://en.wikipedia.org/wiki/Lotka%E2%80%93Volterra_equations$\begin{aligned}{\frac {dx}{dt}} & = \alpha x - \beta xy, \\{\frac {dy}{dt}} & = \delta xy - \gamma y\end{aligned}$where$x$ is the number of preys (for example, rabbits);$y$ is the number of predators (for example, foxes);${\frac {dy}{dt}}$ and ${\frac {dx}{dt}}$ represent the instantaneous growth rates of the two populations;$t$ represents time$α, β, γ, δ$ are positive real parameters describing the interaction of the two species.<jupyter_code># Numpy or eager-mode TF
def lotka_volterra(z, ts, args):
    x, y = z[0], z[1]
    α, β, δ, γ = args[0], args[1], args[2], args[3]
    dx = α*x - β*x*y
    dy = -δ*y + γ*x*y
    return tf.stack([dx, dy], axis=0)

# Tensorflow Variables
D = 2
# time interval from 0 to 10 divided in 100 steps
N = 100
ts = tf.constant(np.linspace(0., 10.0, N))
# initial arguments
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))<jupyter_output><empty_output><jupyter_text>TF-RungeKutta ODE Solver<jupyter_code>sol = ode_tf(lotka_volterra, z0, ts, args)
xs = sol[:, 0].numpy()
ys = sol[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])<jupyter_output><empty_output><jupyter_text>Scipy ODE Solver<jupyter_code>sol = ode_scipy(lotka_volterra, z0, ts, args)
xs = sol[:, 0].numpy()
ys = sol[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])<jupyter_output><empty_output><jupyter_text>The ugly Augmented State in Tensorfow Eager-mode + GradientTapeThis is an implementation of the augmented-state gradient computed by ODE-Solver code sample found in Neural ODE paper made in pure TF 1.x in eager-mode and using Gradient-Tape.To be clear, that was a true pain to implement it and it's still not stable IMHO...I don't pretend this is nice or idiomatic code but it works now and enables experimentations...<jupyter_code>def flatten(params):
    """flatten a list of tensors/arrays into a flat tensor and returns also the unflatten function
    able to rebuild original list from flattenned tensor
    """
    if len(params) == 0:
        def unflatten(flat_params):
            return flat_params
        return np.array(params), unflatten
    
    original_shapes = []
    flat_params = []
    default_dtype = params[0].dtype
    for p in params:
        sh = p.shape
        original_shapes.append(sh)
        flat_sh = 1
        for s in sh:
            flat_sh *= s

        res = tf.reshape(p, (flat_sh,))
        res = tf.cast(res, default_dtype)
        flat_params.append(res)
    
    def unflatten(flat_params):
        cur_idx = 0
        params = []
        for sh in original_shapes:
            flat_sh = 1
            for s in sh:
                flat_sh *= s
            param = flat_params[cur_idx:cur_idx + flat_sh]
            param = tf.reshape(param, sh)
            cur_idx += int(flat_sh)
            params.append(param)
            
        return params
            
    flat_params = tf.concat(flat_params, axis=0)
    return flat_params, unflatten
    
    
# Directly inspired from function implemented with autograd at the end of the paper
# Extended from "Scalable Inference of Ordinary Differential
# Equation Models of Biochemical Processes", Sec. 2.4.2
# <NAME>, <NAME>, <NAME>, 2017
# https://arxiv.org/abs/1711.08079
def grad_odeint(z_ts, func, ts, args, T, D, ode_solver, variables = None):
    """
    This code is still ugly and not very robust IMHO...
    It's full of reshape that might be unneeded but at least I know what I manipulate without compilted types
    
    Args:
        z_ts: the values with ODE Solver in forward path on ts time values
        func: the adjoint gradient function to integrate with ODE Solver
        ts: the discretized time values
        args: the derivable arguments of func
        ode_solver: the ode_solver function (z0, ts, args) -> zs
        variables: the hidden TF/Keras tensor variables (separated from args to be manageable as they have their own life cycle I'm not sure to understand yet)

    Returns:
        grad: the custom augmentate-state TF gradient function
    """
    # let's store
    if variables is not None:
        flat_variables, unflatten_variables = flatten(variables)

    def unpack_D_2D(x, D):
        """simple extract """
        x = tf.squeeze(x)
        return (
            x[0:D],
            x[D:2 * D]
        )
        
    # augmented_state = [z_t, a_z_t, a_t, a_args, a_vars]
    def a_gradient_aug(augmented_state, t, args):
        augmented_state = tf.convert_to_tensor(augmented_state, dtype=z_ts.dtype)
        t = tf.convert_to_tensor(t, dtype=ts.dtype)
        t = tf.reshape(t, (1,))
        args = tf.convert_to_tensor(args, dtype=args.dtype)
        z_t, a_z_t = unpack_D_2D(augmented_state, z_ts.shape[1])

        # the internal gradient-tape computing second derivative
        with tf.GradientTape() as gg:
            gg.watch(args)                
            gg.watch(a_z_t)
            gg.watch(t)
            gg.watch(z_t)
            dz_dt = func(z_t, t, args)

        if variables is not None:
            da_z_dt, da_t_dt, da_args_dargs, da_dvars = gg.gradient(
                dz_dt, [z_t, t, args, variables], output_gradients = -a_z_t)
        else:
            da_z_dt, da_t_dt, da_args_dargs = gg.gradient(
                dz_dt, [z_t, t, args], output_gradients = -a_z_t)
            da_dvars = None
        
        if da_z_dt is None:
            da_z_dt = tf.zeros_like(z_t, dtype=z_t.dtype)
        if da_t_dt is None:
            da_t_dt = tf.zeros_like(t, dtype=t.dtype)
        if da_args_dargs is None:
            da_args_dargs = tf.zeros_like(args, dtype=args.dtype)
        if da_dvars is None:
            da_dvars = tf.zeros_like((1,))
        else:
            da_dvars, _ = flatten(da_dvars)
        r, _ = flatten([dz_dt, da_z_dt, da_t_dt, da_args_dargs, da_dvars])
        return r

    def grad(gt):
        """the TF custom gradient function"""
        
        # initialization of backward gradient computation starting at the end of time interval
        dl_dz_t = gt[-1, :]
        dl_dt0 = tf.constant([0.], shape=(1,), dtype=ts.dtype)
        dl_dargs = tf.zeros_like(args, dtype=ts.dtype)
        if variables is not None:
            dl_dvars = tf.zeros_like(flat_variables, dtype=flat_variables.dtype)
        else:
            dl_dvars = tf.zeros((1,))
        dl_dt_list = []

        # compute gradient back-propagation by ODE solver calls going backward in time
        for i in range(T - 1, 0, -1):
            # Compute effect of moving back in time (just porting formulas found in paper)
            ztsi = z_ts[i, :]
            gti = gt[i, :]
            # Compute the adjoint gradient of loss at current time
            fi = func(ztsi, ts[i], args)
            dl_dt = tf.tensordot(fi, tf.transpose(gti), axes=1) #, axes=[[1], [0]])
            dl_dt = tf.reshape(dl_dt, (1,))
            dl_dt_list.append(dl_dt)
            dl_dt0 = dl_dt0 - dl_dt
            dl_dt0 = tf.reshape(dl_dt0, shape=(1,))

            # build augmented state
            aug_z0, unflatten_aug = flatten([ztsi, dl_dz_t, dl_dt0, dl_dargs, dl_dvars])
            # Run augmented system backwards to the previous observation using ODE_Solver
            cur_ts = tf.reverse(tf.slice(ts, [i-1], [2]), axis=[0])
            a_aug = ode_solver(lambda aug, t, args: a_gradient_aug(aug, t, args), aug_z0, cur_ts, args)

            # get at ts[i - 1]
            a_aug = a_aug[1]
            _, dl_dz_t, dl_dt0, dl_dargs, dl_dvars = unflatten_aug(a_aug)

            # Add gradient from current output.
            dl_dz_t = dl_dz_t + gt[i - 1, :]

        dl_dt0 = tf.reshape(dl_dt0, (1,))
        dl_dt_list.append(dl_dt0)
        dl_dts = tf.stack(dl_dt_list, axis=0)
        dl_dts = tf.reverse(dl_dts, axis=[0])

        # output z_t is not needed but can't put None in TF tensors
        n_o_n_e = tf.zeros_like(z_ts) #, dtype=z_ts.dtype)
        
        if variables is not None:
            return [n_o_n_e, dl_dz_t, dl_dts, dl_dargs], unflatten_variables(dl_dvars)
        else:
            return [n_o_n_e, dl_dz_t, dl_dts, dl_dargs], []
    return grad

def tf_ode_grad(f, N, D, ode_solver):
    """The plug in custom_gradient TF mechanism"""
    @tf.custom_gradient
    def custom_grad(z0, ts, args):
        z0 = tf.squeeze(z0)
        # forward values computation with ODE Solver
        ys = ode_solver(f, z0, ts, args)
        def grad_fn(dr, variables=None):
            # backward gradient computation using backward ODE Solver
            (_, dl_dz_t, dl_dts, dl_dargs), dl_dvars = grad_odeint(ys, f, ts, args, N, D, ode_solver, variables)(dr)
            return [dl_dz_t, dl_dts, dl_dargs], dl_dvars
        return ys, grad_fn
    return custom_grad<jupyter_output><empty_output><jupyter_text>A gradient-descent Neural-ODE optimizer based on loss functionThe following function optimizes iteratively a function defined by its derivative f using an ODE Solver, a ground_truth and a time interval sampling.The derivative function is parameterized by a set of custom arguments and/or a set of hidden Tensorflow variables (the weights & bias in Keras mode for example). Those parameters are those we want to learn during training by gradient descent controlled by an Adam optimizer.I had to separate Tensorflow variables from other arguments because they are managed differently in TF `@custom_gradient`. That was the deepest pain during this study because there is no doc or samples about those features except very basic cases which were useless to me. Anyway, by forcing this separation, I could reach a runnable state that seems to optimize my parameters as you'll see below.<jupyter_code>from IPython import display

def loss_L2(y_true, y_pred):
    """The classic derivable least-square loss function"""
    return tf.reduce_sum(tf.square(y_true - y_pred))

def optimize(f, ode_solver, loss_fn, labels, z0, ts,
             args=None, variables=None,
             epochs=25,
             x_axis_min=-0.25, x_axis_max=10.25,
             y_axis_min=-0.25, y_axis_max=7.0,
             x2_axis_min=-0.25, x2_axis_max=10.25,
             y2_axis_min=-0.25, y2_axis_max=7.0,
             learning_rate=0.1, epsilon=1e-7, output_dir="images", start = 0):
    """A sample optimization process that trie to optimize args/variables parameters of a function
    defined by its derivative f integrated with ODE Solver according to a loss function loss_fn, groundtruth labels,
    initial input value and time interval sampling using Adam optimizer (with plots & image saving)"""
    def custom_grad(z0, ts, the_args):
        with tf.GradientTape() as tape:
            tape.watch(z0)
            tape.watch(the_args)
            ode = tf_ode_grad(f, N, D, ode_solver)(z0, ts, the_args)
            loss = loss_fn(labels, ode)

        # managing stupidly the abscence of args or variables...
        if variables is not None:
            grads, grad_vars = tape.gradient(loss, [[the_args], variables])
        else:
            grads = tape.gradient(loss, [the_args])
            grad_vars = []

        the_grads = []
        for grad in grads:
            grad = tf.reshape(grad, the_args.shape)
            grad = tf.cast(grad, the_args.dtype)
            the_grads.append(grad)
            
        the_grad_vars = []
        for i, grad in enumerate(grad_vars):
            the_grad_vars.append(tf.cast(grad, variables[i].dtype))
            
        return loss, ode, the_grads, the_grad_vars

    # Adam Optimizer
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate,
                                       epsilon=epsilon)

    # plotting layout
    layout = go.Layout(
        xaxis=dict(
            range=[x_axis_min, x_axis_max],
        ),
        yaxis=dict(
            range=[y_axis_min, y_axis_max],
            domain=[0.55, 1.0],
        ),
        xaxis2=dict(
            range=[x2_axis_min, x2_axis_max],
        ),
        yaxis2=dict(
            range=[y2_axis_min, y2_axis_max],
            domain=[0, 0.45],  
        ),
        title="",
        height=680, width=1024
    )
    
    # no args, replacing it by a fake tf.Variable for now... to be removed later
    if args is None:
        var_args = tf.Variable(0, dtype=tf.float64)
    else:
        var_args = args
    
    # the training loop
    global_step = tf.Variable(start)
    best_loss = 1e400
    best_zss = []

    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    
    for epoch in range(epochs):
        loss, ode, grads, grad_vars = custom_grad(z0, ts, var_args)

        print("loss:", loss)
        #print("ode:", ode)
        #print("grads:", grads)

        grad_params = list(zip(grads, [var_args]))
        if variables is not None:
            grad_params.extend(list(zip(grad_vars, variables)))
                    
        optimizer.apply_gradients(grad_params, global_step)
        
        #loss, ode, grads, grad_vars = custom_grad(z0, ts, var_args)
        ys = ode.numpy()
        last_ts = ts.numpy()
        last_ys = ys
        
        if loss < best_loss:
            best_loss = loss
            best_var_args = var_args
            #best_ts = ts.numpy()
            best_zs = ys
            best_zss.append(best_zs)

        # More Plotting
        display.clear_output(wait=True)
        data = []
        
        for zs in best_zss[-1:]:
            data.extend([
                go.Scatter(x=ts.numpy(), y=zs[:, 0], name="best_x",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(34, 119, 158)'),
                                width = 1)),
                go.Scatter(x=ts.numpy(), y=zs[:, 1], name="best_y",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(201, 79, 14)'),
                                width = 1)),
                go.Scatter(x=zs[:, 0], y=zs[:, 1], name="best_xy",
                           xaxis="x2", yaxis="y2",
                           mode = 'lines+markers',
                           line = dict(
                                color = ('rgb(8, 173, 99)'),
                                width = 1)),
            ])
            
        data.extend([            
            go.Scatter(x=ts.numpy(), y=labels.numpy()[:, 0], name="groundtruth_x",
                       line = dict(
                            color = ('rgb(34, 119, 158)'),
                            width = 2,
                            dash = 'dot')),
            go.Scatter(x=ts.numpy(), y=labels.numpy()[:, 1], name="groundtruth_y",
                       line = dict(
                            color = ('rgb(201, 79, 14)'),
                            width = 2,
                            dash = 'dot')),
            go.Scatter(x=labels.numpy()[:, 0], y=labels.numpy()[:, 1], name="groundtruth_xy",
                       xaxis="x2", yaxis="y2",
                       line = dict(
                            color = ('rgb(173, 7, 7)'),
                            width = 2,
                            dash = 'dot')),
        ])
        if loss >= best_loss:
            data.extend([                
                go.Scatter(x=ts.numpy(), y=last_ys[:, 0], name="current_x",
                           line = dict(
                                    color = ('rgb(24, 24, 24)'),
                                    width = 1,
                                    dash = 'dash')),
                go.Scatter(x=ts.numpy(), y=last_ys[:, 1], name="current_y",
                           line = dict(
                                    color = ('rgb(24, 24, 24)'),
                                    width = 1,
                                    dash = 'dash')),
                go.Scatter(x=last_ys[:, 0], y=last_ys[:, 1], name="current_xy",
                           xaxis="x2", yaxis="y2",
                           line = dict(
                                color = ('rgb(24, 24, 24)'),
                                width = 2,
                                dash = 'dash')),
            ])

        layout.title = f"Step {global_step.numpy()}<br>Best Loss {best_loss.numpy()}<br>Current Loss {loss.numpy()}"
        fig = go.Figure(data=data, layout=layout)
        iplot(fig)

        pio.write_image(fig, f"{output_dir}/step_{global_step.numpy()}.png")

        #print(f"Step: {global_step.numpy()}, Loss: {loss.numpy()}")<jupyter_output><empty_output><jupyter_text>Optimizing lotka_volterraLet's optimize lotka_volterra parameters using our optimization function with a dumb loss function that is the squared distance of the output values from 1 to check it works.  Scipy ODE Solver  Model<jupyter_code># Tensorflow Variables
# the dimension of our input vector [x, y] => 2
D = 2
# time interval sampling
N = 25
ts = tf.constant(np.linspace(0., 10.0, N))
# the arguments we want to optimize
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))

def ode_scipy_solver(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

ground_truth_ones = tf.ones(shape=(N, D), dtype=z0.dtype)<jupyter_output><empty_output><jupyter_text>Training Video<jupyter_code>from IPython.display import HTML

HTML("""
<video width="640" height="480" controls>
  <source src="lotka_volterra_optim_scipy/video.mp4" type="video/mp4">
</video>
""")<jupyter_output><empty_output><jupyter_text>Running Optimization by yourself<jupyter_code>optimize(lotka_volterra, ode_scipy_solver, loss_L2,
         ground_truth_ones, z0, ts, args=args, variables=None, epochs=80,
         y_axis_min=-0.25, y_axis_max=7.0,
         x_axis_min=-0.25, x_axis_max=10.25,
         output_dir="lotka_volterra_optim_scipy")<jupyter_output><empty_output><jupyter_text>TF Runge-Kutta ODE Solver (Diverging)  Model<jupyter_code># Tensorflow Variables
# the dimension of our input vector [x, y] => 2
D = 2
# time interval sampling
N = 25
ts = tf.constant(np.linspace(0., 10.0, N))
# the arguments we want to optimize
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))

# the initial input value
z0 = tf.Variable(tf.constant([1.0,1.0], shape=(2,), dtype=ts.dtype))

def ode_tf_solver(f, z0, ts, args):
    ts = tf.contrib.framework.sort(ts)
    r = ode_tf(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

# our ground_truth (all 1.0)
ground_truth_ones = tf.ones(shape=(N, D), dtype=z0.dtype)<jupyter_output><empty_output><jupyter_text>Training Video>Look how it diverges at a first period of optimization.>This is due to computational instability of ODE Solver Runge-Kutta Dopri5 method.>That's why in next parts of this notebook, I'll use on Scipy LSODA implementation that is more much stable.<jupyter_code>from IPython.display import HTML

HTML("""
<video width="640" height="480" controls>
  <source src="lotka_volterra_optim_tf/video.mp4" type="video/mp4">
</video>
""")<jupyter_output><empty_output><jupyter_text>Running Training by yourself<jupyter_code>optimize(lotka_volterra, ode_tf_solver, loss_L2,
         ground_truth_ones, z0, ts, args=args, variables=None, epochs=150,
         y_axis_min=-0.25, y_axis_max=7.0,
         x_axis_min=-0.25, x_axis_max=10.25,
         output_dir="lotka_volterra_optim_tf")<jupyter_output><empty_output><jupyter_text>Final Step: learning time-series with Neural ODEFinally, let's try to learn some time-series function with our neural ODE to check paper assertions are not just imaginary.We are going to model derivative of our function with a classic one-layer dense neural layer as shown in https://julialang.org/blog/2019/01/fluxdiffeq.Then we compute our ground-truth using ODE-Solver from true derivative, decide to use L2 loss.Finally, we launch iterative optimization process using ODE augmented state gradient descent and observe what happens...I won't launch it with Runge-Kutta as it is much less stable than scipy LSODA... tensors are flowing from GPU to CPU but it's OK for current experimentations...  Learning a piece of lotka_volterra curve  Model<jupyter_code>D = 2
N = 20
ts = tf.constant(np.linspace(0., 3.5, N))
args = tf.Variable(tf.constant([1.5, 1.0, 3.0, 1.0], shape=(4,), dtype=ts.dtype))
z0 = tf.Variable(tf.constant([1.0, 1.0], shape=(2,), dtype=ts.dtype))

model = tf.keras.Sequential([
  tf.keras.layers.Dense(25, activation=tf.nn.relu, input_shape=(1, D)),  # input shape required
  tf.keras.layers.Dense(D)
])

def ode_scipy_solver(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

def predict_fn(y, ts, args):
    y2 = tf.reshape(y, [1, D])
    y2 = tf.cast(y2, tf.float32)
    r = model(y2)
    r = tf.squeeze(r)
    r = tf.cast(r, y.dtype)
    return r

ground_truth = ode_scipy(lotka_volterra, z0, ts, args)<jupyter_output><empty_output><jupyter_text>For the lazy, video of trainingRun next cell to see video of a 500 epochs training I did on CPU/GPU (starting at epoch 80 because the rest was about the same).It's quite interesting to see how it learns the curve starting from a kind of average curve and then trying to fit curves progressively...Conclusion: Neural ODE is able to learn a tiny time-series (and certainly overfit a lot on it ;))<jupyter_code>from IPython.display import HTML

HTML("""
<video width="800" height="600" controls>
  <source src="ode_dense_lotka_volterra/video.mp4" type="video/mp4">
</video>
""")<jupyter_output><empty_output><jupyter_text>Running the experiment by yourself> On a 1080TI, it takes around 3h for 1000 epochs<jupyter_code>optimize(predict_fn, ode_scipy_solver, loss_L2, ground_truth, z0, ts,
         args=None, variables=model.variables, epochs=500,
         x_axis_min=-0.25, x_axis_max=3.75,
         y_axis_min=-0.25, y_axis_max=8.0,
         x2_axis_min=-0.25, x2_axis_max=8.0,
         y2_axis_min=-0.25, y2_axis_max=8.0,
         output_dir="ode_dense_lotka_volterra")<jupyter_output><empty_output><jupyter_text>Result of training<jupyter_code>zs = ode_scipy(predict_fn, z0, ts, args)

xs = zs[:, 0].numpy()
ys = zs[:, 1].numpy()

iplot([go.Scatter(x=ts.numpy(), y=xs, name="x"), go.Scatter(x=ts.numpy(), y=ys, name="y")])

iplot([go.Scatter(x=xs, y=ys)])<jupyter_output><empty_output><jupyter_text>Learning a piece of Damped 2D oscillator curve but not converging  Model<jupyter_code>D = 2
N = 50
ts = tf.constant(np.linspace(0., 6.0, N))
z0 = tf.Variable(tf.constant([2.0, 0.0], shape=(2,), dtype=ts.dtype))

args = tf.Variable(tf.constant([
    # −0.1x^3 + 2.0y^3
    [-0.1, 2.0],
    # −2.0x^3 − 0.1y^3
    [-2.0, -0.1]], shape=(2, 2), dtype=ts.dtype))


model2 = tf.keras.Sequential([
  tf.keras.layers.Dense(25, activation=tf.nn.tanh, input_shape=(1, D)),  # input shape required
  tf.keras.layers.Dense(D)
])

def ode_scipy_solver2(f, z0, ts, args):
    r = ode_scipy(f, z0, ts, args)
    r = tf.cast(r, z0.dtype)
    return r

def predict_fn2(y, ts, args):
    y2 = tf.reshape(y, [1, D])
    y2 = tf.cast(y2, tf.float32)
    r = model2(y2)
    r = tf.squeeze(r)
    r = tf.cast(r, y.dtype)
    return r

ground_truth2 = ode_scipy(damped_oscillator_2d, z0, ts, args)<jupyter_output><empty_output><jupyter_text>For the lazy, video of trainingRun next cell to see video of a 1000 epochs training I did on CPU/GPU.It's quite interesting to see how it learns first curves but fails to converge on the whole oscillator.I need to test with a bigger Dense network and maybe another time distribution...Anyway, it's fun to see how it learns, isnt' it? ;)<jupyter_code>from IPython.display import HTML

HTML("""
<video width="800" height="600" controls>
  <source src="ode_dense_damped/video.mp4" type="video/mp4">
</video>
""")<jupyter_output><empty_output><jupyter_text>Running the experiment by yourself> On a 1080TI, it takes around 3h for 1000 epochs<jupyter_code>optimize(predict_fn2, ode_scipy_solver2, loss_L2, ground_truth2, z0, ts,
         args=None, variables=model2.variables, epochs=5,
         x_axis_min=-0.25, x_axis_max=6.25,
         y_axis_min=-2.5, y_axis_max=2.5,
         x2_axis_min=-2.5, x2_axis_max=2.5,
         y2_axis_min=-2.5, y2_axis_max=2.5,
         output_dir="ode_dense_damped_oscillator_2d", start=0)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11290322580645161, 'fraction_numerical': 0.019731819049436914, 'mean_word_length': 3.070408856948626, 'pattern_counts': {'":': 4, '<': 125, '<?xml version=': 0, '>': 160, 'https://': 8, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '25', 'dir': 'jupyter-structured-clean-dedup', 'id': '423441', 'license': 'MIT', 'n_tokens_mistral': 14087, 'n_tokens_neox': 13430, 'n_words': 4423, 'path': 'tf-neural-ode-v1.0.ipynb', 'repo_name': 'mandubian/neural-ode', 'size': '41973'}