I only have the informations of your link. So I dont know about the other issues you talk about.

But if you set for the functional paradigm it is obvious that you need some jax.numpy and that jax.numpy can not implement every numpy function. Numpy and some of its functions (like inplace updates) are inherent non functional. I cant imagine an other way to fix this.