Jacob Kelly

Hi @FedericoV, good catch, thanks for bringing this up! There are probably a number of spots where the docstring is different from what the code actually does, my apologies. If...

I would recommend having a look at the tests for jet which can be found [here](https://github.com/google/jax/blob/master/tests/jet_test.py)

I think `prod` is only available in `python>=3.8`, I just updated the requirements in the README to reflect this. Does that fix your issue?

Hello! Sorry for the delayed reply. I'm having some trouble reproducing this error actually. I used the preprocessed data available in the [release](https://github.com/jacobjinkelly/easy-neural-ode/releases/tag/1.0.0). My conda environment export is: ``` channels:...