jax.Array: support fast path for lax.transpose & lax.squeeze
jax.Array: support fast path for lax.transpose & lax.squeeze

As part of this change, I created a helper function so that the logic of type checking is
in a single location. Eventually we can replace this helper function with appropriate
isinstance() checks using the APIs described in #11859.

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