Replace KeyArray with jax.Array

The private internal alias `KeyArray` was removed from jax in v0.4.36.¹
This commit replaces the private type with the recommended public
type `jax.Array`.²

¹ <jax-ml/jax@fee272e>
² <https://jax.readthedocs.io/en/latest/_autosummary/jax.Array.html#jax.Array>

evofr/infer/MCMC_handler.py

@@ -1,16 +1,15 @@
 import pickle
 from typing import Callable, Dict, Optional, Type
 
-from jax import random
-from jax._src.random import KeyArray
+from jax import random, Array
 from numpyro.infer import MCMC, NUTS, Predictive
 from numpyro.infer.mcmc import MCMCKernel
 
 
 class MCMCHandler:
     def __init__(
         self,
-        rng_key: Optional[KeyArray] = None,
+        rng_key: Optional[Array] = None,
         kernel: Optional[Type[MCMCKernel]] = None,
         **kernel_kwargs
     ):

evofr/infer/SVI_handler.py

@@ -2,8 +2,7 @@
 from typing import Any, Callable, Optional
 
 import jax.example_libraries.optimizers as optimizers
-from jax import random
-from jax._src.random import KeyArray
+from jax import random, Array
 from numpyro.infer import SVI, Predictive, Trace_ELBO
 from numpyro.infer.autoguide import AutoGuide
 from numpyro.infer.svi import SVIState
@@ -14,7 +13,7 @@
 class SVIHandler:
     def __init__(
         self,
-        rng_key: Optional[KeyArray] = None,
+        rng_key: Optional[Array] = None,
         loss: Optional[Trace_ELBO] = None,
         optimizer: Optional[Optimizer] = None,
     ):