dev: metaclass wrap use_state based on a class attr

src/evox/__init__.py

@@ -1,6 +1,5 @@
 from .core.algorithm import Algorithm
 from .core.module import *
-from .core.operator import Operator
 from .core.problem import Problem
 from .core.state import State
 from .core.monitor import Monitor

src/evox/core/algorithm.py

@@ -9,6 +9,7 @@
 
 class Algorithm(Stateful):
     """Base class for all algorithms"""
+    stateful_functions = ['init_ask', 'init_tell', 'ask', 'tell']
 
     def init_ask(self, state: State) -> Tuple[jax.Array, State]:
         """Ask the algorithm for the initial population

src/evox/core/module.py

@@ -129,20 +129,24 @@ def __new__(
         name,
         bases,
         class_dict,
-        force_wrap=["__call__"],
-        ignore=["init", "setup"],
-        ignore_prefix="_",
     ):
-        wrapped = {}
+        wrapped = class_dict
+        stateful_functions = []
+        # stateful_functions from parent classes
+        for base in bases:
+            if hasattr(base, "stateful_functions"):
+                stateful_functions += base.stateful_functions
+
+        # stateful_functions from current class
+        if "stateful_functions" in class_dict:
+            stateful_functions += class_dict["stateful_functions"]
 
         for key, value in class_dict.items():
-            if key in force_wrap:
-                wrapped[key] = use_state(value)
-            elif key.startswith(ignore_prefix) or key in ignore:
-                wrapped[key] = value
-            elif callable(value):
+            if key in stateful_functions:
                 wrapped[key] = use_state(value)
 
+        wrapped["stateful_functions"] = stateful_functions
+
         return super().__new__(cls, name, bases, wrapped)
 
 
@@ -158,6 +162,8 @@ class Stateful(metaclass=MetaStatefulModule):
     and recursively call ``setup`` methods of all submodules.
     """
 
+    stateful_functions = []
+
     def __init__(self) -> None:
         self._node_id = None
         self._cache_override = set()
@@ -201,7 +207,9 @@ def _recursive_init(self, key, node_id, module_name) -> Tuple[State, int]:
                 subkey = None
             else:
                 key, subkey = jax.random.split(key)
-            submodule_state, node_id = attr._recursive_init(subkey, node_id + 1, attr_name)
+            submodule_state, node_id = attr._recursive_init(
+                subkey, node_id + 1, attr_name
+            )
             assert isinstance(
                 submodule_state, State
             ), "setup method must return a State"

src/evox/core/problem.py

@@ -8,6 +8,7 @@
 
 class Problem(Stateful):
     """Base class for all algorithms"""
+    stateful_functions = ["evalutate"]
 
     def evaluate(self, state: State, pop: jax.Array) -> Tuple[jax.Array, State]:
         """Evaluate the fitness at given points

src/evox/core/workflow.py

@@ -0,0 +1,4 @@
+from .module import *
+
+class Workflow(Stateful):
+    stateful_functions = ["step"]

src/evox/operators/selection/find_pbest.py

@@ -1,4 +1,3 @@
-from evox import jit_class, Operator
 from evox.core.state import State
 import jax
 import jax.numpy as jnp