Merge pull request #44 from blab/fix-jax-array-type

Replace `KeyArray` with `jax.Array`
blab · Dec 17, 2024 · 0736267 · 0736267
2 parents 2c91dd8 + 7341ab4
commit 0736267
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/evofr/infer/MCMC_handler.py b/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
     ):

diff --git a/evofr/infer/SVI_handler.py b/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,
     ):