Minor typing issue

[NeilGirdhar]

import jax.numpy as jnp
x: jnp.ndarray = jnp.zeros(10)
y = x.reshape((10,))
reveal_type(x)  # a.py:4: note: Revealed type is 'jax._src.numpy.lax_numpy.ndarray'
reveal_type(y)  # a.py:5: note: Revealed type is 'numpy.ndarray'
x = y  # error: Incompatible types in assignment (expression has type "numpy.ndarray", variable has type "jax._src.numpy.lax_numpy.ndarray")

[jakevdp]

Thanks for the report! I suspect you'll find this issue barely scratches the surface of the typing issues in jax.numpy; almost none of the functions therein are accurately type-annotated, mainly because the effort is blocked by fundamental challenges in applying Python type annotations to JAX arrays (see e.g. #943)

[NeilGirdhar]

@jakevdp Yeah, no worries. I figured it was a ways off. Should I close this issue in favor of #943? I don't want to drown you guys in issues.