Convention for NumPy: onp vs np in Flax

[8bitmp3]

@avital I noticed the convention for ordinary NumPy is onp in some places in google/flax (https://github.com/google/flax/search?q=onp). Should we migrate to np, which is similar to this PR: jax-ml/jax#3760 by @jakevdp? HuggingFace appears to be also using np in its JAX Transformers with Flax.

import numpy as np
import jax.numpy as jnp

Cheers 👍

[avital]

Nice catch -- indeed we should switch away from using onp. Feel free to file an issue or pull request. Thanks!