Neural_ODE
ODE learning using continuous in time back-propagation
Serial_Training
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
Parallel
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
JIT_GPU
Deployment of the NODE using JAX and its JIT module for deployment on CPU, GPU or TPU. Very convenient and good speed up.