Replace `numpy` with `JAX`

**Describe the solution you'd like**
We can't use `numpy` to be quantized because everyone is using frameworks to define their models, so it would make sense to use `JAX` instead as a 2nd backend to be supported.
