brainstate 0.1.0.post20250413__py2.py3-none-any.whl → 0.1.0.post20250422__py2.py3-none-any.whl

brainstate/_compatible_import.py

@@ -0,0 +1,73 @@
+# Copyright 2025 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+import importlib.util
+from contextlib import contextmanager
+from typing import Iterable, Hashable
+
+import jax
+
+__all__ = [
+    'ClosedJaxpr',
+    'Primitive',
+    'extend_axis_env_nd',
+    'jaxpr_as_fun',
+    'get_aval',
+    'Tracer',
+    'to_concrete_aval',
+    'brainevent',
+]
+
+brainevent_installed = importlib.util.find_spec('brainevent') is not None
+
+from jax.core import get_aval, Tracer
+
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import ClosedJaxpr, extend_axis_env_nd, Primitive, jaxpr_as_fun
+else:
+    from jax.extend.core import ClosedJaxpr, Primitive, jaxpr_as_fun
+    from jax.core import trace_ctx
+
+
+    @contextmanager
+    def extend_axis_env_nd(name_size_pairs: Iterable[tuple[Hashable, int]]):
+        prev = trace_ctx.axis_env
+        try:
+            trace_ctx.set_axis_env(prev.extend_pure(name_size_pairs))
+            yield
+        finally:
+            trace_ctx.set_axis_env(prev)
+
+
+def to_concrete_aval(aval):
+    aval = get_aval(aval)
+    if isinstance(aval, Tracer):
+        return aval.to_concrete_value()
+    return aval
+
+
+if brainevent_installed:
+    import brainevent
+else:
+
+    class BrainEvent:
+        def __getattr__(self, item):
+            raise ImportError('brainevent is not installed, please install brainevent first.')
+
+
+    brainevent = BrainEvent()

brainstate/_state.py

@@ -17,12 +17,8 @@ from __future__ import annotations
 
 import contextlib
 import dataclasses
-import jax
-import numpy as np
 import threading
 from functools import wraps, partial
-from jax.api_util import shaped_abstractify
-from jax.extend import source_info_util
 from typing import (
     Any,
     Union,
@@ -39,6 +35,11 @@ from typing import (
     Generator,
 )
 
+import jax
+import numpy as np
+from jax.api_util import shaped_abstractify
+from jax.extend import source_info_util
+
 from brainstate.typing import ArrayLike, PyTree, Missing, Filter
 from brainstate.util import DictManager, PrettyObject
 from brainstate.util.filter import Nothing

brainstate/_state_test.py

@@ -14,9 +14,10 @@
 # ==============================================================================
 
 
-import jax.numpy as jnp
 import unittest
 
+import jax.numpy as jnp
+
 import brainstate as bst
 
 

brainstate/augment/_autograd_test.py

@@ -16,12 +16,13 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
+import unittest
+from pprint import pprint
+
 import brainunit as u
 import jax
 import jax.numpy as jnp
 import pytest
-import unittest
-from pprint import pprint
 
 import brainstate as bst
 from brainstate.augment._autograd import _jacfwd

brainstate/augment/_eval_shape.py

@@ -16,9 +16,10 @@
 from __future__ import annotations
 
 import functools
-import jax
 from typing import Any, TypeVar, Callable, Sequence, Union
 
+import jax
+
 from brainstate import random
 from brainstate.graph import Node, flatten, unflatten
 from ._random import restore_rngs

brainstate/augment/_mapping.py

@@ -955,7 +955,6 @@ def _vmap_new_states_transform(
     if isinstance(axis_size, int) and axis_size <= 0:
         raise ValueError(f"axis_size must be greater than 0, got {axis_size}.")
 
-
     @vmap(
         in_axes=in_axes,
         out_axes=out_axes,

brainstate/augment/_mapping_test.py

@@ -425,6 +425,7 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
 
     def test_incompatible_shapes(self):
         foo = brainstate.nn.LIF(3)
+
         # Simulate an incompatible shapes scenario:
         # We intentionally assign a state with a different shape than expected.
         @bst.augment.vmap_new_states(state_tag='new1', axis_size=5)

brainstate/compile/_ad_checkpoint.py

@@ -16,9 +16,10 @@
 from __future__ import annotations
 
 import functools
-import jax
 from typing import Callable, Tuple, Union
 
+import jax
+
 from brainstate.typing import Missing
 from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
 from ._util import write_back_state_values

brainstate/compile/_conditions.py

@@ -15,11 +15,13 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable, Sequence
+
 import jax
 import jax.numpy as jnp
 import numpy as np
-from collections.abc import Callable, Sequence
 
+from brainstate._compatible_import import to_concrete_aval, Tracer
 from brainstate._utils import set_module_as
 from ._error_if import jit_error_if
 from ._make_jaxpr import StatefulFunction
@@ -86,7 +88,7 @@ def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
             raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
 
     # not jit
-    if jax.config.jax_disable_jit and isinstance(jax.core.get_aval(pred), jax.core.ConcreteArray):
+    if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(pred), Tracer):
         if pred:
             return true_fun(*operands)
         else:

brainstate/compile/_conditions_test.py

@@ -14,9 +14,10 @@
 # ==============================================================================
 from __future__ import annotations
 
+import unittest
+
 import jax
 import jax.numpy as jnp
-import unittest
 
 import brainstate as bst
 

brainstate/compile/_error_if.py

@@ -16,10 +16,11 @@
 from __future__ import annotations
 
 import functools
-import jax
 from functools import partial
 from typing import Callable, Union
 
+import jax
+
 from brainstate._utils import set_module_as
 from ._unvmap import unvmap
 

brainstate/compile/_error_if_test.py

@@ -15,10 +15,11 @@
 
 from __future__ import annotations
 
+import unittest
+
 import jax
 import jax.numpy as jnp
 import jaxlib.xla_extension
-import unittest
 
 import brainstate as bst
 

brainstate/compile/_jit.py

@@ -16,11 +16,12 @@
 from __future__ import annotations
 
 import functools
-import jax
 from collections.abc import Iterable, Sequence
+from typing import (Any, Callable, Union)
+
+import jax
 from jax._src import sharding_impls
 from jax.lib import xla_client as xc
-from typing import (Any, Callable, Union)
 
 from brainstate._utils import set_module_as
 from brainstate.typing import Missing

brainstate/compile/_jit_test.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
-import jax.numpy as jnp
 import unittest
 
+import jax.numpy as jnp
+
 import brainstate as bst
 
 

brainstate/compile/_loop_collect_return.py

@@ -16,11 +16,11 @@
 from __future__ import annotations
 
 import math
+from functools import wraps
+from typing import Callable, Optional, TypeVar, Tuple, Any
 
 import jax
 import jax.numpy as jnp
-from functools import wraps
-from typing import Callable, Optional, TypeVar, Tuple, Any
 
 from brainstate._utils import set_module_as
 from ._make_jaxpr import StatefulFunction

brainstate/compile/_loop_collect_return_test.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
+import unittest
+
 import jax.numpy as jnp
 import numpy as np
-import unittest
 
 import brainstate as bst
 

brainstate/compile/_loop_no_collection.py

@@ -16,9 +16,9 @@
 from __future__ import annotations
 
 import math
+from typing import Any, Callable, TypeVar
 
 import jax
-from typing import Any, Callable, TypeVar
 
 from brainstate._utils import set_module_as
 from ._loop_collect_return import _bounded_while_loop

brainstate/compile/_make_jaxpr.py

@@ -55,10 +55,12 @@ from __future__ import annotations
 
 import functools
 import inspect
-import jax
 import operator
 from collections.abc import Hashable, Iterable, Sequence
 from contextlib import ExitStack
+from typing import Any, Callable, Tuple, Union, Dict, Optional
+
+import jax
 from jax._src import source_info_util
 from jax._src.linear_util import annotate
 from jax._src.traceback_util import api_boundary
@@ -66,18 +68,13 @@ from jax.api_util import shaped_abstractify
 from jax.extend.linear_util import transformation_with_aux, wrap_init
 from jax.interpreters import partial_eval as pe
 from jax.util import wraps
-from typing import Any, Callable, Tuple, Union, Dict, Optional
 
+from brainstate._compatible_import import ClosedJaxpr, extend_axis_env_nd
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
 from brainstate.typing import PyTree
 from brainstate.util import PrettyObject
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import ClosedJaxpr
-else:
-    from jax.extend.core import ClosedJaxpr
-
 AxisName = Hashable
 
 __all__ = [
@@ -200,7 +197,7 @@ class StatefulFunction(PrettyObject):
 
         # implicit parameters
         self.cache_type = cache_type
-        self._cached_jaxpr: Dict[Any, jax.core.ClosedJaxpr] = dict()
+        self._cached_jaxpr: Dict[Any, ClosedJaxpr] = dict()
         self._cached_out_shapes: Dict[Any, PyTree] = dict()
         self._cached_jaxpr_out_tree: Dict[Any, PyTree] = dict()
         self._cached_state_trace: Dict[Any, StateTraceStack] = dict()
@@ -210,7 +207,7 @@ class StatefulFunction(PrettyObject):
             return None
         return k, v
 
-    def get_jaxpr(self, cache_key: Hashable = ()) -> jax.core.ClosedJaxpr:
+    def get_jaxpr(self, cache_key: Hashable = ()) -> ClosedJaxpr:
         """
         Read the JAX Jaxpr representation of the function.
 
@@ -507,8 +504,8 @@ def make_jaxpr(
     return_shape: bool = False,
     abstracted_axes: Optional[Any] = None,
     state_returns: Union[str, Tuple[str, ...]] = ('read', 'write')
-) -> Callable[..., (Tuple[jax.core.ClosedJaxpr, Tuple[State, ...]] |
-                    Tuple[jax.core.ClosedJaxpr, Tuple[State, ...], PyTree])]:
+) -> Callable[..., (Tuple[ClosedJaxpr, Tuple[State, ...]] |
+                    Tuple[ClosedJaxpr, Tuple[State, ...], PyTree])]:
     """
     Creates a function that produces its jaxpr given example args.
 
@@ -754,12 +751,12 @@ def _make_jaxpr(
             debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
         with ExitStack() as stack:
             if axis_env is not None:
-                stack.enter_context(jax.core.extend_axis_env_nd(axis_env))
+                stack.enter_context(extend_axis_env_nd(axis_env))
             if jax.__version_info__ < (0, 5, 0):
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
             else:
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)
-        closed_jaxpr = jax.core.ClosedJaxpr(jaxpr, consts)
+        closed_jaxpr = ClosedJaxpr(jaxpr, consts)
         if return_shape:
             out_avals, _ = jax.util.unzip2(out_type)
             out_shapes_flat = [jax.ShapeDtypeStruct(a.shape, a.dtype) for a in out_avals]

brainstate/compile/_make_jaxpr_test.py

@@ -15,17 +15,14 @@
 
 from __future__ import annotations
 
+import unittest
+
 import jax
 import jax.numpy as jnp
 import pytest
-import unittest
 
 import brainstate as bst
-
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import jaxpr_as_fun
-else:
-    from jax.extend.core import jaxpr_as_fun
+from brainstate._compatible_import import jaxpr_as_fun
 
 
 class TestMakeJaxpr(unittest.TestCase):

brainstate/compile/_progress_bar.py

@@ -17,9 +17,10 @@ from __future__ import annotations
 
 import copy
 import importlib.util
-import jax
 from typing import Optional, Callable, Any, Tuple, Dict
 
+import jax
+
 tqdm_installed = importlib.util.find_spec('tqdm') is not None
 
 __all__ = [

brainstate/compile/_unvmap.py

@@ -20,13 +20,9 @@ import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
 
+from brainstate._compatible_import import Primitive
 from brainstate._utils import set_module_as
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import Primitive
-else:
-    from jax.extend.core import Primitive
-
 __all__ = [
     "unvmap",
 ]

brainstate/environ.py

@@ -17,18 +17,18 @@
 
 from __future__ import annotations
 
-from collections import defaultdict
-
 import contextlib
 import dataclasses
 import functools
-import numpy as np
 import os
 import re
 import threading
+from collections import defaultdict
+from typing import Any, Callable, Dict, Hashable
+
+import numpy as np
 from jax import config, devices, numpy as jnp
 from jax.typing import DTypeLike
-from typing import Any, Callable, Dict, Hashable
 
 from .mixin import Mode
 

brainstate/environ_test.py

@@ -14,9 +14,10 @@
 # ==============================================================================
 
 
-import jax.numpy as jnp
 import unittest
 
+import jax.numpy as jnp
+
 import brainstate as bst
 
 

brainstate/functional/_activations.py

@@ -20,10 +20,11 @@ Shared neural network activations and other functions.
 
 from __future__ import annotations
 
+from typing import Any, Union, Sequence
+
 import brainunit as u
 import jax
 from jax.scipy.special import logsumexp
-from typing import Any, Union, Sequence
 
 from brainstate import random
 from brainstate.typing import ArrayLike

brainstate/functional/_activations_test.py

@@ -16,12 +16,12 @@
 """Tests for nn module."""
 
 import itertools
+from functools import partial
 
 import jax
 import jax.numpy as jnp
 import scipy.stats
 from absl.testing import parameterized
-from functools import partial
 from jax._src import test_util as jtu
 from jax.test_util import check_grads
 

brainstate/functional/_normalization.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
+from typing import Optional, Union
+
 import brainunit as u
 import jax
-from typing import Optional, Union
 
 from brainstate._utils import set_module_as
 from brainstate.typing import ArrayLike

brainstate/functional/_others.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
+from functools import partial
+
 import jax
 import jax.numpy as jnp
-from functools import partial
 
 from brainstate.typing import PyTree
 

brainstate/graph/_graph_operation.py

@@ -18,10 +18,11 @@
 from __future__ import annotations
 
 import dataclasses
-import jax
-import numpy as np
 from typing import (Any, Callable, Generic, Iterable, Iterator, Mapping, MutableMapping,
                     Sequence, Type, TypeVar, Union, Hashable, Tuple, Dict, Optional, overload)
+
+import jax
+import numpy as np
 from typing_extensions import TypeGuard, Unpack
 
 from brainstate._state import State, TreefyState

brainstate/graph/_graph_operation_test.py

@@ -15,13 +15,14 @@
 
 from __future__ import annotations
 
-import jax
-import jax.numpy as jnp
 import unittest
-from absl.testing import absltest, parameterized
 from collections.abc import Callable
 from threading import Thread
 
+import jax
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
+
 import brainstate as bst
 
 

brainstate/init/_base.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
-import numpy as np
 from typing import Optional, Tuple
 
+import numpy as np
+
 from brainstate.util import PrettyRepr, PrettyType, PrettyAttr
 
 __all__ = ['Initializer', 'to_size']

brainstate/init/_generic.py

@@ -16,10 +16,11 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
+from typing import Union, Callable, Optional, Sequence
+
 import brainunit as bu
 import jax
 import numpy as np
-from typing import Union, Callable, Optional, Sequence
 
 from brainstate._state import State
 from brainstate._utils import set_module_as

brainstate/nn/__init__.py

@@ -25,6 +25,8 @@ from ._dynamics import *
 from ._dynamics import __all__ as dynamics_all
 from ._elementwise import *
 from ._elementwise import __all__ as elementwise_all
+from ._event import *
+from ._event import __all__ as _event_all
 from ._exp_euler import *
 from ._exp_euler import __all__ as exp_euler_all
 from ._interaction import *
@@ -45,6 +47,7 @@ __all__ = (
     + exp_euler_all
     + interaction_all
     + utils_all
+    + _event_all
 )
 
 del (
@@ -57,4 +60,5 @@ del (
     exp_euler_all,
     interaction_all,
     utils_all,
+    _event_all,
 )

brainstate/nn/_collective_ops.py

@@ -19,6 +19,7 @@ from collections import namedtuple
 from typing import Callable, TypeVar, Tuple, Any, Dict
 
 import jax
+from typing import Callable, TypeVar, Tuple, Any, Dict
 
 from brainstate._state import catch_new_states
 from brainstate._utils import set_module_as

brainstate/nn/_collective_ops_test.py

@@ -41,7 +41,3 @@ class Test_init_all_states:
         gru = bst.nn.GRUCell(1, 2)
         bst.nn.init_all_states(gru, batch_size=10)
         print(gru)
-
-
-
-

brainstate/nn/_common.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 from collections import defaultdict
-
 from typing import Any, Sequence, Hashable, Dict
 
 from brainstate import environ

brainstate/nn/_dyn_impl/__init__.py

@@ -20,8 +20,6 @@ from ._dynamics_synapse import *
 from ._dynamics_synapse import __all__ as dyn_synapse_all
 from ._inputs import *
 from ._inputs import __all__ as inputs_all
-from ._projection_alignpost import *
-from ._projection_alignpost import __all__ as alignpost_all
 from ._rate_rnns import *
 from ._rate_rnns import __all__ as rate_rnns
 from ._readout import *
@@ -31,7 +29,6 @@ __all__ = (
     dyn_neuron_all
     + dyn_synapse_all
     + inputs_all
-    + alignpost_all
     + rate_rnns
     + readout_all
 )
@@ -41,6 +38,5 @@ del (
     dyn_synapse_all,
     inputs_all,
     readout_all,
-    alignpost_all,
     rate_rnns,
 )