refactor: util routine

wcxve · Jan 31, 2024 · f4a2b61 · f4a2b61
1 parent a45fea4
commit f4a2b61
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/elisa/__init__.py b/src/elisa/__init__.py
@@ -1,5 +1,10 @@
-from .model import *  # noqa: F403
+from . import data, infer, model, util
+from .data import *  # noqa F403
+from .infer import *  # noqa F403
+from .model import *  # noqa F403
 from .util import jax_enable_x64, set_cpu_cores
 
+__all__ = data.__all__ + infer.__all__ + model.__all__ + util.__all__
+
 jax_enable_x64(True)
 set_cpu_cores(4)
diff --git a/src/elisa/util/__init__.py b/src/elisa/util/__init__.py
@@ -1 +1,3 @@
-from .config import *  # noqa: F403
+from .config import jax_enable_x64, set_cpu_cores
+
+__all__ = ['jax_enable_x64', 'set_cpu_cores']