ODE learning using continuous in time back-propagation
Uses autograd
to implement the Neural-ODE algorithm Chen, Tian Qi, et al. "Neural ordinary differential equations." Advances in neural information processing systems. 2018. in serial
Uses autograd
as well as mpi4py
to to run parallel trainings of the neural ODE with gradient information exchange at each epoch (will add a conditional statement to allow for update after a preset number of epochs) - implemented for a different time series
Deployment of the NODE using JAX and its JIT module for deployment on CPU, GPU or TPU. Very convenient and good speed up.