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

brainstate/_compatible_import.py

@@ -16,6 +16,7 @@
 # -*- coding: utf-8 -*-
 
 
+import importlib.util
 from contextlib import contextmanager
 from typing import Iterable, Hashable
 
@@ -29,8 +30,11 @@ __all__ = [
     '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):
@@ -56,3 +60,14 @@ def to_concrete_aval(aval):
         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,17 +15,17 @@
 
 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
 from ._util import wrap_single_fun_in_multi_branches, write_back_state_values
-from brainstate._compatible_import import to_concrete_aval, Tracer
-
 
 __all__ = [
     'cond', 'switch', 'ifelse',

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

@@ -57,7 +57,7 @@ import functools
 import inspect
 import operator
 from collections.abc import Hashable, Iterable, Sequence
-from contextlib import ExitStack, contextmanager
+from contextlib import ExitStack
 from typing import Any, Callable, Tuple, Union, Dict, Optional
 
 import jax
@@ -69,11 +69,11 @@ 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 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
-from brainstate._compatible_import import ClosedJaxpr, extend_axis_env_nd
 
 AxisName = Hashable
 

brainstate/compile/_make_jaxpr_test.py

@@ -15,10 +15,11 @@
 
 from __future__ import annotations
 
+import unittest
+
 import jax
 import jax.numpy as jnp
 import pytest
-import unittest
 
 import brainstate as bst
 from brainstate._compatible_import import jaxpr_as_fun

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,9 +20,8 @@ import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
 
-from brainstate._utils import set_module_as
 from brainstate._compatible_import import Primitive
-
+from brainstate._utils import set_module_as
 
 __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,
 )