neural network dynamics on jax

Pipeline:

1. Simulate teacher network using pre_training_routine
2. Run odes on pre training routine results
3. Move results from cluster to local, or run directly on cluster
4. Run pre-process plot data
5. Plot data in plot_cluster.ipynb