switch PRNGKey to key

Author: younik

chromax/functional.py

@@ -25,7 +25,7 @@ def cross(
     :param recombination_vec: array of m probabilities.
         The i-th value represent the probability to recombine before the marker i.
     :type recombination_vec:
-    :param random_key: JAX Array with dtype jax.dtypes.prng_key, for reproducibility purpose.
+    :param random_key: JAX random key, for reproducibility purpose.
     :type random_key: jax.Array 
     :param mutation_probability: The probability of having a mutation in a marker.
     :type mutation_probability: float
@@ -43,7 +43,7 @@ def cross(
         >>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
         >>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
         >>> rec_vec = rec_vec.flatten()
-        >>> random_key = jax.random.PRNGKey(42)
+        >>> random_key = jax.random.key(42)
         >>> f2 = functional.cross(parents, rec_vec, random_key)
         >>> f2.shape
         (50, 1000, 2)
@@ -52,7 +52,7 @@ def cross(
     random_keys = jax.random.split(
         random_key, num=2 * len(parents) * 2 * parents.shape[3]
     )
-    random_keys = random_keys.reshape(2, len(parents), 2, parents.shape[3], 2)
+    random_keys = random_keys.reshape(2, len(parents), 2, parents.shape[3])
     cross_random_key, mutate_random_key = random_keys
 
     offsprings = _cross(
@@ -100,8 +100,8 @@ def double_haploid(
     :param recombination_vec: array of m probabilities.
         The i-th value represent the probability to recombine before the marker i.
     :type recombination_vec: ndarray
-    :param random_key: array of n PRNGKey, one for each individual.
-    :type random_key: jax.Array with dtype jax.dtypes.prng_key
+    :param random_key: JAX random key, for reproducibility purpose.
+    :type random_key: jax.Array
     :param mutation_probability: The probability of having a mutation in a marker.
     :type mutation_probability: float
     :return: output population of shape (n, n_offspring, m, d).
@@ -118,15 +118,15 @@ def double_haploid(
         >>> rec_vec = np.full((n_chr, chr_len), 1.5 / chr_len)
         >>> rec_vec[:, 0] = 0.5  # equal probability on starting haploid
         >>> rec_vec = rec_vec.flatten()
-        >>> random_key = jax.random.PRNGKey(42)
+        >>> random_key = jax.random.key(42)
         >>> dh = functional.double_haploid(f1, 10, rec_vec, random_key)
         >>> dh.shape
         (50, 10, 1000, 2)
     """
     population = population.reshape(*population.shape[:2], -1, 2)
     keys = jax.random.split(
         random_key, num=2 * len(population) * n_offspring * population.shape[2]
-    ).reshape(2, len(population), n_offspring, population.shape[2], 2)
+    ).reshape(2, len(population), n_offspring, population.shape[2])
     cross_random_key, mutate_random_key = keys
     haploids = _double_haploid(
         population,

chromax/simulator.py

@@ -171,7 +171,7 @@ def set_seed(self, seed: int):
         :param seed: random seed.
         :type seed: int
         """
-        self.random_key = jax.random.PRNGKey(seed)
+        self.random_key = jax.random.key(seed)
 
     def load_population(self, file_name: Union[Path, str]) -> Population["n"]:
         """Load a population from file.
@@ -271,7 +271,7 @@ def differentiable_cross_func(self) -> Callable:
             >>> f1 = simulator.load_population(sample_data.genome)
             >>> weights = np.random.uniform(size=(10, len(f1), 2))
             >>> weights /= weights.sum(axis=1, keepdims=True)
-            >>> random_key = jax.random.PRNGKey(42)
+            >>> random_key = jax.random.key(42)
             >>> grad_value = grad_f(f1, weights, random_key)
             >>> grad_value.shape
             (10, 371, 2)
@@ -284,7 +284,7 @@ def differentiable_cross_func(self) -> Callable:
         def diff_cross_f(
             population: Population["n"],
             cross_weights: Float[Array, "m n 2"],
-            random_key: jax.random.PRNGKeyArray,
+            random_key: jax.Array,
         ) -> Population["m"]:
             population = population.reshape(*population.shape[:-1], -1, 2)
             keys_shape = len(cross_weights), len(population), 2, population.shape[-2]

pyproject.toml

@@ -21,7 +21,7 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     'Intended Audience :: Science/Research',
 ]
-dependencies = ["numpy", "pandas", "jax", "jaxlib", "jaxtyping"]
+dependencies = ["numpy", "pandas", "jax>=0.4.16", "jaxlib>=0.4.16", "jaxtyping"]
 dynamic = ["version"]
 
 [project.optional-dependencies]

tests/test_functional.py

@@ -15,7 +15,7 @@ def test_cross(idx):
     parents = np.random.choice([False, True], size=parents_shape)
     rec_vec = np.zeros(n_markers)
     rec_vec[0] = idx
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     new_pop = functional.cross(parents, rec_vec, random_key)
 
     for i in range(ploidy):
@@ -28,7 +28,7 @@ def test_double_haploid():
     pop_shape = (50, n_chr * chr_len, ploidy)
     f1 = np.random.choice([False, True], size=pop_shape)
     rec_vec = np.full((n_chr * chr_len,), 1.5 / chr_len)
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     dh = functional.double_haploid(f1, n_offspring, rec_vec, random_key)
     assert dh.shape == (len(f1), n_offspring, n_chr * chr_len, ploidy)
 
@@ -62,7 +62,7 @@ def test_cross_mutation():
     rec_vec = np.full((n_markers,), 1.5e-2)
     cross = functional.cross
 
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     assert np.all(cross(zeros_pop, rec_vec, random_key) == 0)
     assert np.all(cross(zeros_pop, rec_vec, random_key, 1) == 1)
     mutated_pop = cross(zeros_pop, rec_vec, random_key, 0.5)
@@ -83,7 +83,7 @@ def test_dh_mutation():
     rec_vec = np.full((n_markers,), 1.5e-2)
     dh = functional.double_haploid
 
-    random_key = jax.random.PRNGKey(42)
+    random_key = jax.random.key(42)
     assert np.all(dh(zeros_pop, 10, rec_vec, random_key) == 0)
     assert np.all(dh(zeros_pop, 10, rec_vec, random_key, 1) == 1)
     mutated_pop = dh(zeros_pop, 10, rec_vec, random_key, 0.5)