brainstate 0.1.1__py2.py3-none-any.whl → 0.1.3__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.3"
 
 from . import augment
 from . import compile

brainstate/_compatible_import.py

@@ -19,7 +19,7 @@
 import importlib.util
 from contextlib import contextmanager
 from functools import partial
-from typing import Iterable, Hashable, TypeVar, Callable
+from typing import Iterable, Hashable, TypeVar, Callable, TYPE_CHECKING
 
 import jax
 
@@ -37,6 +37,7 @@ __all__ = [
     'unzip2',
     'wraps',
     'Device',
+    'wrap_init',
 ]
 
 T = TypeVar("T")
@@ -44,6 +45,8 @@ T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 
+from saiunit._compatible_import import wrap_init
+
 brainevent_installed = importlib.util.find_spec('brainevent') is not None
 
 from jax.core import get_aval, Tracer
@@ -148,13 +151,13 @@ def to_concrete_aval(aval):
     return aval
 
 
-if brainevent_installed:
-    import brainevent
-else:
-
-    class BrainEvent:
-        def __getattr__(self, item):
-            raise ImportError('brainevent is not installed, please install brainevent first.')
+if not brainevent_installed:
+    if not TYPE_CHECKING:
+        class BrainEvent:
+            def __getattr__(self, item):
+                raise ImportError('brainevent is not installed, please install brainevent first.')
 
+        brainevent = BrainEvent()
 
-    brainevent = BrainEvent()
+else:
+    import brainevent

brainstate/_state.py

@@ -1049,7 +1049,7 @@ class StateTraceStack(Generic[A]):
         """
         if self._jax_trace_new_arg is not None:
             # internal use
-            state._value = jax.tree.map(lambda x: self._jax_trace_new_arg(shaped_abstractify(x)), state._value)
+            state._value = jax.tree.map(self._jax_trace_new_arg, state._value)
 
     def __enter__(self) -> 'StateTraceStack':
         TRACE_CONTEXT.state_stack.append(self)