numpy
jax
diffmah