Array vs PyTree anotation about jaxtyping HOT 2 CLOSED

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.

from jaxtyping.

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

from jaxtyping.