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.

[jakevdp]

The design doc is coming together here: #11859

In particular, it does answer the main question of this Issue

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

The answer is No, because we need inputs to functions to be far more flexible than ndarray. Inputs should be ArrayLike, where this accepts jax.numpy.ndarray, numpy.ndarray, numpy scalars, Python scalars, objects with __array__ attributes, etc.

For return values, we can annotate with jnp.ndarray, or perhaps a new jax.typing.NDArray.

[NeilGirdhar]

The answer is No, because we need inputs to functions to be far more flexible than ndarray. Inputs should be ArrayLike, where this accepts jax.numpy.ndarray, numpy.ndarray, numpy scalars, Python scalars, objects with array attributes, etc.

Would it be possible to create a type alias for jax.numpy.typing.ArrayLike like the one in Numpy?

[YouJiacheng]

actually in jax, the problem is simpler than numpy, since jax.numpy is not designed to be compatible with sequence, e.g. list. It is not too bad to use an union of scalar, ndarray and supportsjaxarray.

[jakevdp]

Would it be possible to create a type alias for jax.numpy.typing.ArrayLike like the one in Numpy?

Yes, the draft doc I linked to mentions doing that in the roadmap.

[jakevdp]

I'm going to close this issue, because I think the work going forward is better outlined and tracked elsewhere. If you have thoughts about the type annotation roadmap, please feel free to comment on #11859

[jakevdp]

Hi @YouJiacheng - if you're still interested in this, we've started the process of doing these sorts of annotations in #12049. See the attached PRs for a few examples. Sorry it took such a long time to get there, but we'd love your help now if you'd like to contribute!