[typing] add class-level declarations of Array members.
[typing] add class-level declarations of Array members.

This fixes some pytype errors associated with the changes in #12421

google/jax,google/jax,google/jax,google/jax