Use onp instead of np in ode_test (#3288)
Use onp instead of np in ode_test (#3288)

* Use onp instead of np in ode_test

* other ode_test.py fixes

Co-authored-by: Matthew Johnson <e68678f07030b810d1440f7f1c1c423f5afe56fe@google.com>
google/jax,google/jax,google/jax,tensorflow/probability,google/jax,tensorflow/probability