wip(hsgp): switch typing to jax.Array

should probably do this everywhere, add reference to jax typing user guide

numpyro/contrib/hsgp/util.py

@@ -1,6 +1,7 @@
 from typing import Union
 
+import jax
 from jaxlib.xla_extension import ArrayImpl
 import numpy as np
 
-ARRAY_TYPE = Union[ArrayImpl, np.ndarray]
+ARRAY_TYPE = Union[jax.Array, np.ndarray]  # jax.Array covers tracers