Add comment clarifying the intent of the Array/ArrayLike annotations

Author: jakevdp

jax/_src/typing.py

@@ -45,13 +45,16 @@ class HasDtypeAttribute(Protocol):
 DimSize = Union[int, Any]  # extensible
 Shape = Sequence[DimSize]
 
-# Array is a type annotation for JAX arrays and tracers.
+# Array is a type annotation for standard JAX arrays and tracers produced by
+# core functions in jax.lax and jax.numpy; it is not meant to include
+# future non-standard array types like KeyArray and BInt.
 # For now we set it to Any; in the future this will be more restrictive
 # (see https://github.com/google/jax/pull/11859)
 # TODO(jakevdp): make this conform to the JEP 12049 plan.
 Array = Any
 
-# ArrayLike is a Union of all objects that can be implicitly converted to a JAX array.
+# ArrayLike is a Union of all objects that can be implicitly converted to a standard
+# JAX array (i.e. not including future non-standard array types like KeyArray and BInt).
 ArrayLike = Union[
   Array,  # JAX array type
   np.ndarray,  # NumPy array type