dev: add parallel_init and parallel_step as a sugar

src/evox/core/module.py

@@ -269,6 +269,30 @@ def init(self, key: jax.Array = None, no_state: bool = False) -> State:
         state, _node_id = self._recursive_init(key, 0, None, no_state)
         return state
 
+    def parallel_init(
+        self, key: jax.Array, num_copies: int, no_state: bool = False
+    ) -> Tuple[State, int]:
+        """Initialize multiple copies of this module in parallel
+
+        This method should not be overwritten.
+
+        Parameters
+        ----------
+        key
+            A PRNGKey.
+        num_copies
+            The number of copies to be initialized
+        no_state
+            Whether to skip the state initialization
+
+        Returns
+        -------
+        Tuple[State, int]
+            The state of this module and all submodules combined, and the last node_id
+        """
+        subkeys = jax.random.split(key, num_copies)
+        return jax.vmap(self.init, in_axes=(0, None))(subkeys, no_state)
+
     @classmethod
     def stack(cls, stateful_objs, axis=0):
         for obj in stateful_objs:

src/evox/workflows/std_workflow.py

@@ -1,5 +1,6 @@
 from collections.abc import Callable, Sequence
-from typing import Optional, Union
+from typing import Optional, Tuple, Union
+from functools import partial
 import warnings
 
 import jax
@@ -23,6 +24,13 @@
 from evox.utils import parse_opt_direction
 
 
+def _leftover_callbacks_warning(method_name):
+    warnings.warn(
+        f"`{method_name}` is called with a state that has leftover callbacks. "
+        "Did you forget to call `execute_callbacks`?"
+    )
+
+
 @dataclass
 class StdWorkflowState:
     generation: int
@@ -189,11 +197,12 @@ def _step(self, state):
 
             return state
 
+        self._step = partial(_step, self)
+        self._parallel_step = jax.vmap(self._step)
         if self.jit_step:
             # the first argument is self, which should be static
-            self._step = jax.jit(_step, static_argnums=(0,))
-        else:
-            self._step = _step
+            self._step = jax.jit(self._step)
+            self._parallel_step = jax.jit(self._parallel_step)
 
         # by default, use the first device
         self.devices = jax.local_devices()[:1]
@@ -292,12 +301,19 @@ def setup(self, key):
 
     def step(self, state):
         if self.auto_exec_callbacks and state._callbacks:
-            warnings.warn(
-                "`step` is called with a state that has leftover callbacks."
-                "Did you forget to call `execute_callbacks`?"
-            )
+            _leftover_callbacks_warning("step")
+
+        state = self._step(state)
+
+        if self.auto_exec_callbacks:
+            state = state.execute_callbacks(state)
+        return state
+
+    def parallel_step(self, state):
+        if self.auto_exec_callbacks and state._callbacks:
+            _leftover_callbacks_warning("parallel_step")
 
-        state = self._step(self, state)
+        state = self._parallel_step(state)
 
         if self.auto_exec_callbacks:
             state = state.execute_callbacks(state)