last updates to address PR reviews

Author: Lisa

examples/aurora.ipynb

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

examples/cmaes.ipynb

@@ -208,18 +208,21 @@
     "covs = [(state.sigma**2) * state.cov_matrix]\n",
     "\n",
     "iteration_count = 0\n",
+    "sample_fn = jax.jit(cmaes.sample)\n",
+    "update_fn = jax.jit(cmaes.update)\n",
+    "stop_condition_fn = jax.jit(cmaes.stop_condition)\n",
     "for _ in range(num_iterations):\n",
     "    iteration_count += 1\n",
     "\n",
     "    # sample\n",
     "    key, subkey = jax.random.split(key)\n",
-    "    samples = jax.jit(cmaes.sample)(state, subkey)\n",
+    "    samples = sample_fn(state, subkey)\n",
     "\n",
     "    # update\n",
-    "    state = jax.jit(cmaes.update)(state, samples)\n",
+    "    state = update_fn(state, samples)\n",
     "\n",
     "    # check stop condition\n",
-    "    stop_condition = jax.jit(cmaes.stop_condition)(state)\n",
+    "    stop_condition = stop_condition_fn(state)\n",
     "\n",
     "    if stop_condition:\n",
     "        break\n",

examples/mapelites_asktell.ipynb

@@ -327,20 +327,23 @@
     "except ImportError:\n",
     "    bar = range(num_iterations)\n",
     "\n",
+    "ask_fn = jax.jit(map_elites.ask)\n",
+    "tell_fn = jax.jit(map_elites.tell)\n",
+    "\n",
     "# Main loop\n",
     "for i in bar:\n",
     "    start_time = time.time()\n",
     "    key, subkey = jax.random.split(key)\n",
     "    # Generate solutions\n",
-    "    genotypes, extra_info = map_elites.ask(repertoire, emitter_state, subkey)\n",
+    "    genotypes, extra_info = ask_fn(repertoire, emitter_state, subkey)\n",
     "\n",
     "    # Evaluate solutions: get fitness, descriptor and extra scores.\n",
     "    # This is where custom evaluations on CPU or GPU can be added.\n",
     "    key, subkey = jax.random.split(key)\n",
     "    fitnesses, descriptors, extra_scores = scoring_fn(genotypes, subkey)\n",
     "\n",
     "    # Update MAP-Elites\n",
-    "    repertoire, emitter_state, current_metrics = map_elites.tell(\n",
+    "    repertoire, emitter_state, current_metrics = tell_fn(\n",
     "        genotypes=genotypes,\n",
     "        fitnesses=fitnesses,\n",
     "        descriptors=descriptors,\n",

examples/mees.ipynb

@@ -448,11 +448,8 @@
    "provenance": []
   },
   "gpuClass": "standard",
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -466,7 +463,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

examples/nsga2_spea2.ipynb

@@ -296,7 +296,6 @@
     "\n",
     "# init spea2\n",
     "key, subkey = jax.random.split(key)\n",
-    "init_fn = partial(spea2.init, population_size=population_size)\n",
     "repertoire, emitter_state, init_metrics = spea2.init(\n",
     "    genotypes,\n",
     "    population_size,\n",

examples/pga_aurora.ipynb

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

qdax/core/containers/archive.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from functools import partial
 from typing import Any, Dict, Tuple
 
 import jax
@@ -240,8 +239,7 @@ def score_euclidean_novelty(
     Returns:
         The novelty scores of the given state descriptors.
     """
-    knn_fn = partial(knn, k=num_nearest_neighb)
-    values, _indices = knn_fn(archive.data, state_descriptors)
+    values, _indices = knn(archive.data, state_descriptors, num_nearest_neighb)
 
     summed_distances = jnp.mean(jnp.square(values), axis=1)
     return scaling_ratio * summed_distances
@@ -278,8 +276,7 @@ def knn(
     dist = jnp.sqrt(jnp.clip(dist, min=0.0))
 
     # return values, indices
-    qdax_top_k_fn = partial(qdax_top_k, k=k)
-    values, indices = qdax_top_k_fn(-dist)
+    values, indices = qdax_top_k(-dist, k=k)
 
     return -values, indices
 

qdax/utils/plotting.py

@@ -1,6 +1,5 @@
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-import jax
 import jax.numpy as jnp
 import matplotlib as mpl
 import matplotlib.cm as cm
@@ -456,9 +455,7 @@ def plot_mome_pareto_fronts(
     axes[1].set_ylim(minval, maxval)
 
     if with_global:
-        global_pareto_front, pareto_bool = jax.jit(
-            repertoire.compute_global_pareto_front
-        )()
+        global_pareto_front, pareto_bool = repertoire.compute_global_pareto_front()
         global_pareto_descriptors = jnp.concatenate(repertoire_descriptors)[pareto_bool]
         axes[0].scatter(
             global_pareto_front[:, 0],

tests/core_test/aurora_test.py

@@ -218,13 +218,15 @@ def observation_extractor_fn(
     target_repertoire_size = 1024
 
     previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - target_repertoire_size
+    update_fn = jax.jit(aurora.update)
+    container_size_control_fn = jax.jit(aurora.container_size_control)
 
     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, _ = jax.jit(aurora.update)(
+            repertoire, emitter_state, _ = update_fn(
                 repertoire,
                 emitter_state,
                 subkey,
@@ -244,7 +246,7 @@ def observation_extractor_fn(
 
         elif iteration % 2 == 0:
             # only CSC
-            repertoire, previous_error = jax.jit(aurora.container_size_control)(
+            repertoire, previous_error = container_size_control_fn(
                 repertoire,
                 target_size=target_repertoire_size,
                 previous_error=previous_error,
@@ -427,22 +429,23 @@ def observation_extractor_fn(
     target_repertoire_size = 1024
 
     previous_error = jnp.sum(repertoire.fitnesses != -jnp.inf) - target_repertoire_size
+    ask_fn = jax.jit(aurora.ask)
+    tell_fn = jax.jit(aurora.tell)
+    container_size_control_fn = jax.jit(aurora.container_size_control)
 
     iteration = 0
 
     while iteration < max_iterations:
         # standard MAP-Elites-like loop
         for _ in range(log_freq):
             key, subkey = jax.random.split(key)
-            genotypes, extra_info = jax.jit(aurora.ask)(
-                repertoire, emitter_state, subkey
-            )
+            genotypes, extra_info = ask_fn(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, _ = jax.jit(aurora.tell)(
+            repertoire, emitter_state, _ = tell_fn(
                 genotypes=genotypes,
                 fitnesses=fitnesses,
                 descriptors=descriptors,
@@ -466,7 +469,7 @@ def observation_extractor_fn(
 
         elif iteration % 2 == 0:
             # only CSC
-            repertoire, previous_error = jax.jit(aurora.container_size_control)(
+            repertoire, previous_error = container_size_control_fn(
                 repertoire,
                 target_size=target_repertoire_size,
                 previous_error=previous_error,

tests/core_test/map_elites_test.py

@@ -269,22 +269,22 @@ def play_step_fn(
         key=key,
         extra_scores=extra_scores,
     )
+    ask_fn = jax.jit(map_elites.ask)
+    tell_fn = jax.jit(map_elites.tell)
 
     # Run the algorithm
     for _ in range(num_iterations):
         key, subkey = jax.random.split(key)
         # Generate solutions
-        genotypes, extra_info = jax.jit(map_elites.ask)(
-            repertoire, emitter_state, subkey
-        )
+        genotypes, extra_info = ask_fn(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 = jax.jit(map_elites.tell)(
+        repertoire, emitter_state, current_metrics = tell_fn(
             genotypes=genotypes,
             fitnesses=fitnesses,
             descriptors=descriptors,