Can we replace `Array = Any` with `Array = ndarray` now?

[YouJiacheng]

It seems that problems in #943 have been resolved?

[jakevdp]

I don't think it's as simple as that – for example, inputs to JAX functions will accept any array-like object, including Python scalar types, numpy scalar types, numpy arrays, jax arrays/tracers, plus for many functions, arbitrary objects that define a __jax_array__ method. So we'd need to account for those too, which makes it more complicated than simply setting Array = jnp.ndarray.

[jakevdp]

I personally prefer Array = Any because its false positive rate is zero 😁

[YouJiacheng]

Oh, I forget __jax_array__.
However, IIUC, we can still annotate the return type as ndarray, and its false positive rate is zero as well.
And return type annotation is very helpful for IDE auto-completing and semantic highlighting

[YouJiacheng]

@jakevdp WDYT about annotate the return type as ndarray, or even Union[ndarray, Any] (IIUC Union[ndarray, PRNGKeyArray] is enough), which can be very useful for development with IDE/vscode.

[jakevdp]

Sure, we could try it, and just be ready to roll it back quickly if it breaks things.

[jakevdp]

I have to say though, I have a knee-jerk dislike for using Union[Any, ndarray] when the result is an array. What it tells me is that the static typing system is not actually expressive enough to statically type our API, and if that's necessary I'd prefer to continue to just use Any or leave things undeclared.

[YouJiacheng]

@jakevdp
I have an idea: use overload to provide a fully annotated version for the most usual case, and a Any annotated version for corner case.

[NeilGirdhar]

What's Jax's long term plan with this? Perhaps the Array API will provide annotations that can be used in Jax?

[jakevdp]

I wrote a design doc a while ago regarding how the project should approach type annotations (surprise, it's not that simple because "type annotations" are a multiply-overloaded concept in Python!)... it got lost in the shuffle. I'll dust it off and try to generate some discussion.

[YouJiacheng]

❤️Does the design doc also discuss things mentioned in #10322 ?

[YouJiacheng]

BTW, now I know __jax_array__ can be solved by Protocol, i.e. Union[ndarray, JAXArray] where JAXArray is a protocol implemented __jax_array__(self) -> ndarray.

[jakevdp]

It's still unclear whether __jax_array__ is something we want to support broadly; see the discussion in #10065.