Array vs PyTree anotation

[gerdm]

Hi!

I'm curious why the API for array annotations is Float[Array, "dims"] and not Array[float, "dims"]? The latter would make it consistent with PyTree[float].

[patrick-kidger]

So from a static type-checking point of view, we'd like both Array and Array[...] to mean jnp.ndarray.

We could accomplish the former by just setting Array = jnp.ndarray, but then the latter would break ([...] isn't valid: you can't write jnp.ndarray[...].

The fix is to do what we do here. As far as the static type checker is concerned, Float = typing.Annotated, and thus Float[Array, ...] = Array = jnp.ndarray, and things work out.

TL;DR: limitations of static type checking.

[gerdm]

Thanks for the explanation @patrick-kidger! That makes sense.