neural network dynamics on jax
Pipeline:
- Simulate teacher network using pre_training_routine
- Run odes on pre training routine results
- Move results from cluster to local, or run directly on cluster
- Run pre-process plot data
- Plot data in plot_cluster.ipynb