rnp: Optionally execute your JAX code with a NumPy backend.

[planes]

Hey Folks! My colleagues and I have always struggled with moving our scientific computing NumPy code bases over to JAX. There's a handful of reasons that people have given us for not migrating to JAX.

1. NumPy code is typically significantly faster for a single iteration on a CPU. While JAX jit-ed GPU code is faster for multiple iterations (optimization/dataset building). Blasphemous, but not all applications need gradients.
2. Unfortunately, many of our collaborators in the engineering world have IT restrictions that make it hard for them to use JAX.
3. We have students that don't have the latest and greatest machines available, so having simpler ecosystems is helpful.
4. We have found that it's easier to debug and verify our code in NumPy.

However, the benefits of JAX are too great to pass up. Our developers work in a world that blends scientific computing and AI. It was hard to pick just JAX or just NumPy and rather than try to support dual code bases, instead we made a very thin wrapper of both. I apologize if this has already been done, but we couldn't find a package that does this already.

rnp (Research NumPy) allows one to code using JAX syntax and with one line switch back ends from JAX to NumPy and back.

We believe this provides the best of both worlds. The single inference speed of NumPy and the power of JAX when you need it. We foresee there being follow on use cases beyond scientific computing, such as local CPU AI inference.

Feedback, issues, and pull requests are welcome.

Thanks!