fix tests

Author: Lisa

examples/pga_aurora.ipynb

@@ -514,7 +514,6 @@
     "n_target = 1024\n",
     "\n",
     "previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - n_target\n",
-    "container_size_control_jitted = jax.jit(aurora.container_size_control)\n",
     "\n",
     "iteration = 0\n",
     "while iteration < max_iterations:\n",
@@ -543,7 +542,7 @@
     "        )\n",
     "\n",
     "    elif iteration % 2 == 0:\n",
-    "        repertoire, previous_error = container_size_control_jitted(\n",
+    "        repertoire, previous_error = jax.jit(aurora.container_size_control)(\n",
     "            repertoire,\n",
     "            target_size=n_target,\n",
     "            previous_error=previous_error,\n",

tests/core_test/aurora_test.py

@@ -218,15 +218,13 @@ def observation_extractor_fn(
     target_repertoire_size = 1024
 
     previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - target_repertoire_size
-    container_size_control_jitted = jax.jit(aurora.container_size_control)
-    update_jitted = jax.jit(aurora.update)
 
     iteration = 0
     while iteration < max_iterations:
         # standard MAP-Elites-like loop
         for _ in range(log_freq):
             key, subkey = jax.random.split(key)
-            repertoire, emitter_state, _ = update_jitted(
+            repertoire, emitter_state, _ = jax.jit(aurora.update)(
                 repertoire,
                 emitter_state,
                 subkey,
@@ -246,7 +244,7 @@ def observation_extractor_fn(
 
         elif iteration % 2 == 0:
             # only CSC
-            repertoire, previous_error = container_size_control_jitted(
+            repertoire, previous_error = jax.jit(aurora.container_size_control)(
                 repertoire,
                 target_size=target_repertoire_size,
                 previous_error=previous_error,
@@ -431,21 +429,20 @@ def observation_extractor_fn(
     previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - target_repertoire_size
 
     iteration = 0
-    container_size_control_jitted = jax.jit(aurora.container_size_control)
-    ask_jitted = jax.jit(aurora.ask)
-    tell_jitted = jax.jit(aurora.tell)
 
     while iteration < max_iterations:
         # standard MAP-Elites-like loop
         for _ in range(log_freq):
             key, subkey = jax.random.split(key)
-            genotypes, extra_info = ask_jitted(repertoire, emitter_state, subkey)
+            genotypes, extra_info = jax.jit(aurora.ask)(
+                repertoire, emitter_state, subkey
+            )
 
             # scores the offsprings
             key, subkey = jax.random.split(key)
             fitnesses, descriptors, extra_scores = aurora_scoring_fn(genotypes, subkey)
 
-            repertoire, emitter_state, _ = tell_jitted(
+            repertoire, emitter_state, _ = jax.jit(aurora.tell)(
                 genotypes=genotypes,
                 fitnesses=fitnesses,
                 descriptors=descriptors,
@@ -469,7 +466,7 @@ def observation_extractor_fn(
 
         elif iteration % 2 == 0:
             # only CSC
-            repertoire, previous_error = container_size_control_jitted(
+            repertoire, previous_error = jax.jit(aurora.container_size_control)(
                 repertoire,
                 target_size=target_repertoire_size,
                 previous_error=previous_error,

tests/core_test/map_elites_test.py

@@ -270,22 +270,21 @@ def play_step_fn(
         extra_scores=extra_scores,
     )
 
-    ask_jitted = jax.jit(map_elites.ask)
-    tell_jitted = jax.jit(map_elites.tell)
-
     # Run the algorithm
     for _ in range(num_iterations):
         key, subkey = jax.random.split(key)
         # Generate solutions
-        genotypes, extra_info = ask_jitted(repertoire, emitter_state, subkey)
+        genotypes, extra_info = jax.jit(map_elites.ask)(
+            repertoire, emitter_state, subkey
+        )
 
         # Evaluate solutions: get fitness, descriptor and extra scores.
         # This is where custom evaluations on CPU or GPU can be added.
         key, subkey = jax.random.split(key)
         fitnesses, descriptors, extra_scores = scoring_fn(genotypes, subkey)
 
         # Update MAP-Elites
-        repertoire, emitter_state, current_metrics = tell_jitted(
+        repertoire, emitter_state, current_metrics = jax.jit(map_elites.tell)(
             genotypes=genotypes,
             fitnesses=fitnesses,
             descriptors=descriptors,