Tracking Issue: JAX Type Annotations

[jakevdp]

This is an issue to track progress on JAX Type Annotations.

For information, see the Type Annotation Roadmap, added in #11859