Rename all arguments named "unused" by "_"

examples/aurora.ipynb

@@ -322,7 +322,7 @@
     "centroids = jnp.zeros(shape=(num_centroids, aurora_dims))\n",
     "\n",
     "@jax.jit\n",
-    "def update_scan_fn(carry: Any, unused: Any) -> Any:\n",
+    "def update_scan_fn(carry: Any, _: Any) -> Any:\n",
     "    \"\"\"Scan the update function.\"\"\"\n",
     "    repertoire, key, aurora_extra_info = carry\n",
     "\n",

examples/pga_aurora.ipynb

@@ -368,7 +368,7 @@
     "centroids = jnp.zeros(shape=(num_centroids, aurora_dims))\n",
     "\n",
     "@jax.jit\n",
-    "def update_scan_fn(carry: Any, unused: Any) -> Any:\n",
+    "def update_scan_fn(carry: Any, _: Any) -> Any:\n",
     "    \"\"\"Scan the update function.\"\"\"\n",
     "    (\n",
     "        repertoire,\n",

qdax/baselines/genetic_algorithm.py

@@ -131,15 +131,15 @@ def update(
     def scan_update(
         self,
         carry: Tuple[GARepertoire, Optional[EmitterState], RNGKey],
-        unused: Any,
+        _: Any,
     ) -> Tuple[Tuple[GARepertoire, Optional[EmitterState], RNGKey], Metrics]:
         """Rewrites the update function in a way that makes it compatible with the
         jax.lax.scan primitive.
 
         Args:
             carry: a tuple containing the repertoire, the emitter state and a
                 random key.
-            unused: unused element, necessary to respect jax.lax.scan API.
+            _: unused element, necessary to respect jax.lax.scan API.
 
         Returns:
             The updated repertoire and emitter state, with a new random key and metrics.

qdax/core/containers/archive.py

@@ -311,7 +311,7 @@ def top_1(data: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
         return data, value, indice
 
     def scannable_top_1(
-        carry: jnp.ndarray, unused: Any
+        carry: jnp.ndarray, _: Any
     ) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]:
         data = carry
         data, value, indice = top_1(data)

qdax/core/distributed_map_elites.py

@@ -191,7 +191,7 @@ def get_distributed_update_fn(
         @jax.jit
         def _scan_update(
             carry: Tuple[MapElitesRepertoire, Optional[EmitterState], RNGKey],
-            unused: Any,
+            _: Any,
         ) -> Tuple[Tuple[MapElitesRepertoire, Optional[EmitterState], RNGKey], Metrics]:
             """Rewrites the update function in a way that makes it compatible with the
             jax.lax.scan primitive."""

qdax/core/emitters/cma_pool_emitter.py

@@ -45,7 +45,7 @@ def batch_size(self) -> int:
         Returns:
             the batch size emitted by the emitter.
         """
-        return self._emitter.batch_size
+        return self._emitter.batch_size  # type: ignore
 
     @partial(jax.jit, static_argnames=("self",))
     def init(
@@ -69,9 +69,7 @@ def init(
             The initial state of the emitter.
         """
 
-        def scan_emitter_init(
-            carry: RNGKey, unused: Any
-        ) -> Tuple[RNGKey, CMAEmitterState]:
+        def scan_emitter_init(carry: RNGKey, _: Any) -> Tuple[RNGKey, CMAEmitterState]:
             key = carry
             key, subkey = jax.random.split(key)
             emitter_state = self._emitter.init(

qdax/core/emitters/qpg_emitter.py

@@ -312,7 +312,7 @@ def state_update(
         emitter_state = emitter_state.replace(replay_buffer=replay_buffer)
 
         def scan_train_critics(
-            carry: QualityPGEmitterState, unused: Any
+            carry: QualityPGEmitterState, _: Any
         ) -> Tuple[QualityPGEmitterState, Any]:
             emitter_state = carry
             new_emitter_state = self._train_critics(emitter_state)
@@ -501,7 +501,7 @@ def _mutation_function_pg(
 
         def scan_train_policy(
             carry: Tuple[QualityPGEmitterState, Genotype, optax.OptState],
-            unused: Any,
+            _: Any,
         ) -> Tuple[Tuple[QualityPGEmitterState, Genotype, optax.OptState], Any]:
             emitter_state, policy_params, policy_optimizer_state = carry
             (

qdax/core/map_elites.py

@@ -154,15 +154,15 @@ def update(
     def scan_update(
         self,
         carry: Tuple[MapElitesRepertoire, Optional[EmitterState], RNGKey],
-        unused: Any,
+        _: Any,
     ) -> Tuple[Tuple[MapElitesRepertoire, Optional[EmitterState], RNGKey], Metrics]:
         """Rewrites the update function in a way that makes it compatible with the
         jax.lax.scan primitive.
 
         Args:
             carry: a tuple containing the repertoire, the emitter state and a
                 random key.
-            unused: unused element, necessary to respect jax.lax.scan API.
+            _: unused element, necessary to respect jax.lax.scan API.
 
         Returns:
             The updated repertoire and emitter state, with a new random key and metrics.

qdax/utils/uncertainty_metrics.py

@@ -288,7 +288,7 @@ def _perform_reevaluation(
 
         def _sampling_scan(
             key: RNGKey,
-            unused: Tuple[()],
+            _: Tuple[()],
         ) -> Tuple[Tuple[RNGKey], Tuple[Fitness, Descriptor, ExtraScores]]:
             key, subkey = jax.random.split(key)
             (

tests/baselines_test/mees_test.py

@@ -163,7 +163,7 @@ def metrics_function(repertoire: MapElitesRepertoire) -> Dict:
     repertoire, emitter_state = map_elites.init(init_variables, centroids, subkey)
 
     @jax.jit
-    def update_scan_fn(carry: Any, unused: Any) -> Any:
+    def update_scan_fn(carry: Any, _: Any) -> Any:
         # iterate over grid
         repertoire, emitter_state, key = carry
         key, subkey = jax.random.split(key)

tests/baselines_test/pgame_test.py

@@ -177,7 +177,7 @@ def metrics_function(repertoire: MapElitesRepertoire) -> Dict:
     repertoire, emitter_state = map_elites.init(init_variables, centroids, subkey)
 
     @jax.jit
-    def update_scan_fn(carry: Any, unused: Any) -> Any:
+    def update_scan_fn(carry: Any, _: Any) -> Any:
         # iterate over grid
         repertoire, emitter_state, key = carry
         key, subkey = jax.random.split(key)

tests/baselines_test/qdpg_test.py

@@ -227,7 +227,7 @@ def metrics_function(repertoire: MapElitesRepertoire) -> Dict:
     repertoire, emitter_state = map_elites.init(init_variables, centroids, subkey)
 
     @jax.jit
-    def update_scan_fn(carry: Any, unused: Any) -> Any:
+    def update_scan_fn(carry: Any, _: Any) -> Any:
         # iterate over grid
         repertoire, emitter_state, key = carry
         key, subkey = jax.random.split(key)