replace onp, np with np, jnp

[benjaminpope]

@shashankdholakia has strong feelings that we should replace onp, np with np, jnp. I tend to agree. Opening this issue in case @LouisDesdoigts has strong feelings otherwise.

[LouisDesdoigts]

I tend to prefer the former, since I never really use normal numpy at all. From what I have seen of most jax code-bases import as np, unless there is large parts of the code that use normal numpy.

I would say we should probably just use whatever convention the other jax packages in the ecosystem near us use, but ultimately its preference.