Put __all__ in complex modules to hide internal functions

Author: BillHuang2001

docs/source/conf.py

@@ -48,7 +48,9 @@
 ]
 
 autodoc2_module_all_regexes = [
-    r"evox\..*",
+    r"evox\.core\..*",
+    r"evox\.problems\..*",
+    r"evox\.workflows\..*",
 ]
 
 autodoc2_render_plugin = "myst"

src/evox/core/components.py

@@ -1,3 +1,11 @@
+__all__ = [
+    "Algorithm",
+    "Problem",
+    "Workflow",
+    "Monitor",
+]
+
+
 from abc import ABC
 from typing import Any, Dict
 

src/evox/core/module.py

@@ -1,3 +1,14 @@
+__all__ = [
+    "Parameter",
+    "Mutable",
+    "ModuleBase",
+    "TransformGetSetItemToIndex",
+    "compile",
+    "vmap",
+    "use_state",
+]
+
+
 from functools import wraps
 from typing import Callable, Dict, Optional, TypeVar, Union
 
@@ -34,9 +45,7 @@ def Parameter(
     )
 
 
-def Mutable(
-    value: torch.Tensor, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None
-) -> torch.Tensor:
+def Mutable(value: torch.Tensor, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None) -> torch.Tensor:
     """Wraps a value as a mutable tensor.
     This is often used to label a value in an algorithm as a mutable tensor that may changes during iteration(s).
 
@@ -172,7 +181,7 @@ def use_state(stateful_func: Union[Callable, nn.Module]) -> Callable:
     def wrapper(params_and_buffers: Dict[str, torch.Tensor], *args, **kwargs):
         params_and_buffers = {("_inner_module." + k): v for k, v in params_and_buffers.items()}
         output = torch.func.functional_call(module, params_and_buffers, *args, **kwargs)
-        params_and_buffers = {k[len("_inner_module."):]: v for k, v in params_and_buffers.items()}
+        params_and_buffers = {k[len("_inner_module.") :]: v for k, v in params_and_buffers.items()}
         if output is None:
             return params_and_buffers
         else:

src/evox/problems/__init__.py

@@ -1,4 +1,4 @@
-__all__ = ["neuroevolution", "numerical"]
+__all__ = ["neuroevolution", "numerical", "hpo_wrapper"]
 
 
-from . import neuroevolution, numerical
+from . import neuroevolution, numerical, hpo_wrapper

src/evox/problems/hpo_wrapper.py

@@ -1,3 +1,11 @@
+__all__ = [
+    "HPOMonitor",
+    "HPOFitnessMonitor",
+    "HPOProblemWrapper",
+    "HPOData",
+]
+
+
 import weakref
 from abc import ABC
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple

src/evox/problems/neuroevolution/__init__.py

@@ -0,0 +1,13 @@
+__all__ = [
+    "brax",
+    "mujoco_playground",
+    "supervised_learning",
+    "BraxProblem",
+    "MujocoPlaygroundProblem",
+    "SupervisedLearningProblem",
+]
+
+from . import brax, mujoco_playground, supervised_learning
+from .brax import BraxProblem
+from .mujoco_playground import MujocoPlaygroundProblem
+from .supervised_learning import SupervisedLearningProblem

src/evox/problems/neuroevolution/utils.py

@@ -1,3 +1,9 @@
+__all__ = [
+    "ModelStateForwardResult",
+    "get_vmap_model_state_forward",
+]
+
+
 import copy
 from typing import Callable, Dict, NamedTuple, Tuple
 

src/evox/problems/numerical/__init__.py

@@ -1,4 +1,7 @@
 __all__ = [
+    "basic",
+    "cec2022",
+    "dtlz",
     "Ackley",
     "Griewank",
     "Rastrigin",
@@ -23,6 +26,7 @@
     "ellipsoid_func",
 ]
 
+from . import basic, cec2022, dtlz
 from .basic import (
     Ackley,
     Ellipsoid,

src/evox/problems/numerical/basic.py

@@ -1,3 +1,22 @@
+__all__ = [
+    "ShiftAffineNumericalProblem",
+    "Ackley",
+    "Griewank",
+    "Rastrigin",
+    "Rosenbrock",
+    "Schwefel",
+    "Sphere",
+    "Ellipsoid",
+    "ackley_func",
+    "griewank_func",
+    "rastrigin_func",
+    "rosenbrock_func",
+    "schwefel_func",
+    "sphere_func",
+    "ellipsoid_func",
+]
+
+
 import torch
 
 from evox.core import Problem

src/evox/problems/numerical/cec2022.py

@@ -1,3 +1,8 @@
+__all__ = [
+    "CEC2022",
+]
+
+
 import os
 from math import ceil
 from typing import List, Optional