Merge pull request #423 from JaxGaussianProcesses/update_keyarray_typing

Update typing.py
JaxGaussianProcesses · Nov 29, 2023 · ac47576 · ac47576
2 parents 48f2db9 + 33e3ce1
commit ac47576
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/gpjax/typing.py b/gpjax/typing.py
@@ -17,17 +17,18 @@
     Callable,
     Union,
 )
-from jax.random import KeyArray as JAXKeyArray
 from jaxtyping import (
     Array as JAXArray,
     Bool,
     Float,
     Int,
+    Key,
     UInt32,
 )
 from numpy import ndarray as NumpyArray
 
 OldKeyArray = UInt32[JAXArray, "2"]
+JAXKeyArray = Key[JAXArray, ""]
 KeyArray = Union[
     OldKeyArray, JAXKeyArray
 ]  # for compatibility regardless of enable_custom_prng setting