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

brainstate/nn/_exp_euler.py

@@ -16,9 +16,10 @@
 
 from __future__ import annotations
 
+from typing import Callable
+
 import brainunit as u
 import jax.numpy as jnp
-from typing import Callable
 
 from brainstate import environ, random
 from brainstate.augment import vector_grad

brainstate/nn/_exp_euler_test.py

@@ -15,9 +15,10 @@
 
 from __future__ import annotations
 
-import brainunit as u
 import unittest
 
+import brainunit as u
+
 import brainstate as bst
 
 

brainstate/nn/_interaction/_conv.py

@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 import collections.abc
+from typing import Callable, Tuple, Union, Sequence, Optional, TypeVar
+
 import jax
 import jax.numpy as jnp
-from typing import Callable, Tuple, Union, Sequence, Optional, TypeVar
 
 from brainstate import init, functional
 from brainstate._state import ParamState

brainstate/nn/_interaction/_linear.py

@@ -17,9 +17,10 @@
 
 from __future__ import annotations
 
+from typing import Callable, Union, Optional
+
 import brainunit as u
 import jax.numpy as jnp
-from typing import Callable, Union, Optional
 
 from brainstate import init, functional
 from brainstate._state import ParamState

brainstate/nn/_interaction/_linear_test.py

@@ -16,8 +16,9 @@
 
 from __future__ import annotations
 
-import brainunit as u
 import unittest
+
+import brainunit as u
 from absl.testing import parameterized
 
 import brainstate as bst

brainstate/nn/_interaction/_normalizations.py

@@ -17,9 +17,10 @@
 
 from __future__ import annotations
 
+from typing import Callable, Union, Sequence, Optional, Any
+
 import jax
 import jax.numpy as jnp
-from typing import Callable, Union, Sequence, Optional, Any
 
 from brainstate import environ, init
 from brainstate._state import ParamState, BatchState

brainstate/nn/_interaction/_poolings.py

@@ -17,13 +17,14 @@
 
 from __future__ import annotations
 
-import brainunit as u
 import functools
+from typing import Sequence, Optional
+from typing import Union, Tuple, Callable, List
+
+import brainunit as u
 import jax
 import jax.numpy as jnp
 import numpy as np
-from typing import Sequence, Optional
-from typing import Union, Tuple, Callable, List
 
 from brainstate import environ
 from brainstate.nn._module import Module

brainstate/nn/_module_test.py

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

brainstate/nn/metrics.py

@@ -16,13 +16,14 @@
 
 from __future__ import annotations
 
-import jax
-import jax.numpy as jnp
-import numpy as np
 import typing as tp
 from dataclasses import dataclass
 from functools import partial
 
+import jax
+import jax.numpy as jnp
+import numpy as np
+
 from brainstate._state import State
 
 __all__ = [

brainstate/optim/_lr_scheduler.py

@@ -16,10 +16,11 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
+from typing import Sequence, Union
+
 import jax
 import jax.numpy as jnp
 import numpy as np
-from typing import Sequence, Union
 
 from brainstate import environ
 from brainstate._state import State, LongTermState

brainstate/optim/_lr_scheduler_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/optim/_optax_optimizer_test.py

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

brainstate/optim/_sgd_optimizer.py

@@ -16,11 +16,12 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-import brainunit as u
 import functools
+from typing import Union, Dict, Optional, Tuple, Any, TypeVar
+
+import brainunit as u
 import jax
 import jax.numpy as jnp
-from typing import Union, Dict, Optional, Tuple, Any, TypeVar
 
 from brainstate import environ
 from brainstate._state import State, LongTermState, StateDictManager

brainstate/random/_rand_funs.py

@@ -17,9 +17,10 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-import numpy as np
 from typing import Optional
 
+import numpy as np
+
 from brainstate.typing import DTypeLike, Size, SeedOrKey
 from ._rand_state import RandomState, DEFAULT
 

brainstate/random/_rand_funs_test.py

@@ -15,12 +15,13 @@
 
 from __future__ import annotations
 
+import platform
+import unittest
+
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
-import platform
 import pytest
-import unittest
 
 import brainstate as bst
 

brainstate/random/_rand_seed.py

@@ -14,11 +14,12 @@
 # ==============================================================================
 from __future__ import annotations
 
-import jax
-import numpy as np
 from contextlib import contextmanager
 from typing import Optional
 
+import jax
+import numpy as np
+
 from brainstate.typing import SeedOrKey
 from ._rand_state import RandomState, DEFAULT, use_prng_key
 

brainstate/random/_rand_seed_test.py

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

brainstate/random/_rand_state.py

@@ -16,16 +16,17 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
+from functools import partial
+from operator import index
+from typing import Optional
+
 import brainunit as u
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
-from functools import partial
 from jax import jit, vmap
 from jax import lax, core, dtypes
-from operator import index
-from typing import Optional
 
 from brainstate import environ
 from brainstate._state import State

brainstate/surrogate.py

@@ -21,13 +21,9 @@ import jax.numpy as jnp
 import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
+from brainstate._compatible_import import Primitive
 from brainstate.util._pretty_pytree import PrettyObject
 
-if jax.__version_info__ < (0, 4, 38):
-    from jax.core import Primitive
-else:
-    from jax.extend.core import Primitive
-
 __all__ = [
     'Surrogate',
     'Sigmoid',

brainstate/typing.py

@@ -16,18 +16,18 @@
 from __future__ import annotations
 
 import builtins
-
-import brainunit as u
 import functools as ft
 import importlib
 import inspect
-import jax
-import numpy as np
 from typing import (
     Any, Callable, Hashable, List, Protocol, Tuple, TypeVar, Union,
     runtime_checkable, TYPE_CHECKING, Generic, Sequence
 )
 
+import brainunit as u
+import jax
+import numpy as np
+
 tp = importlib.import_module("typing")
 
 __all__ = [

brainstate/util/_caller.py

@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 import dataclasses
-import jax
 from typing import Any, TypeVar, Protocol, Generic
 
+import jax
+
 __all__ = [
     'DelayedAccessor',
     'CallableProxy',

brainstate/util/_others.py

@@ -15,17 +15,17 @@
 
 from __future__ import annotations
 
-import gc
-
 import copy
 import functools
-import jax
+import gc
 import threading
 import types
 from collections.abc import Iterable
-from jax.lib import xla_bridge
 from typing import Any, Callable, Tuple, Union, Dict
 
+import jax
+from jax.lib import xla_bridge
+
 from brainstate._utils import set_module_as
 
 __all__ = [

brainstate/util/_pretty_pytree.py

@@ -18,9 +18,9 @@
 from __future__ import annotations
 
 from collections import abc
+from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict
 
 import jax
-from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict
 
 from brainstate.typing import Filter, PathParts
 from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object

brainstate/util/_pretty_pytree_test.py

@@ -15,8 +15,9 @@
 
 from __future__ import annotations
 
-import jax
 import unittest
+
+import jax
 from absl.testing import absltest
 
 import brainstate as bst

brainstate/util/_pretty_table.py

@@ -422,14 +422,14 @@ class PrettyTable:
     def _column_specific_args(self):
         # Column specific arguments, use property.setters
         for attr in (
-            "align",
-            "valign",
-            "max_width",
-            "min_width",
-            "int_format",
-            "float_format",
-            "custom_format",
-            "none_format",
+                "align",
+                "valign",
+                "max_width",
+                "min_width",
+                "int_format",
+                "float_format",
+                "custom_format",
+                "none_format",
         ):
             setattr(
                 self, attr, (self._kwargs[attr] or {}) if attr in self._kwargs else {}
@@ -517,15 +517,15 @@ class PrettyTable:
         elif option == "none_format":
             self._validate_none_format(val)
         elif option in (
-            "start",
-            "end",
-            "max_width",
-            "min_width",
-            "min_table_width",
-            "max_table_width",
-            "padding_width",
-            "left_padding_width",
-            "right_padding_width",
+                "start",
+                "end",
+                "max_width",
+                "min_width",
+                "min_table_width",
+                "max_table_width",
+                "padding_width",
+                "left_padding_width",
+                "right_padding_width",
         ):
             self._validate_nonnegative_int(option, val)
         elif option == "sortby":
@@ -539,16 +539,16 @@ class PrettyTable:
         elif option == "fields":
             self._validate_all_field_names(option, val)
         elif option in (
-            "header",
-            "border",
-            "preserve_internal_border",
-            "reversesort",
-            "xhtml",
-            "format",
-            "print_empty",
-            "oldsortslice",
-            "escape_header",
-            "escape_data",
+                "header",
+                "border",
+                "preserve_internal_border",
+                "reversesort",
+                "xhtml",
+                "format",
+                "print_empty",
+                "oldsortslice",
+                "escape_header",
+                "escape_data",
         ):
             self._validate_true_or_false(option, val)
         elif option == "header_style":
@@ -561,18 +561,18 @@ class PrettyTable:
             for k, formatter in val.items():
                 self._validate_function(f"{option}.{k}", formatter)
         elif option in (
-            "vertical_char",
-            "horizontal_char",
-            "horizontal_align_char",
-            "junction_char",
-            "top_junction_char",
-            "bottom_junction_char",
-            "right_junction_char",
-            "left_junction_char",
-            "top_right_junction_char",
-            "top_left_junction_char",
-            "bottom_right_junction_char",
-            "bottom_left_junction_char",
+                "vertical_char",
+                "horizontal_char",
+                "horizontal_align_char",
+                "junction_char",
+                "top_junction_char",
+                "bottom_junction_char",
+                "right_junction_char",
+                "left_junction_char",
+                "top_right_junction_char",
+                "top_left_junction_char",
+                "bottom_right_junction_char",
+                "bottom_left_junction_char",
         ):
             self._validate_single_char(option, val)
         elif option == "attributes":
@@ -1998,8 +1998,8 @@ class PrettyTable:
         if options["header"]:
             lines.append(self._stringify_header(options))
         elif options["border"] and options["hrules"] in (
-            HRuleStyle.ALL,
-            HRuleStyle.FRAME,
+                HRuleStyle.ALL,
+                HRuleStyle.FRAME,
         ):
             lines.append(self._stringify_hrule(options, where="top_"))
             if title and options["vrules"] in (VRuleStyle.ALL, VRuleStyle.FRAME):
@@ -2115,8 +2115,8 @@ class PrettyTable:
             if options["hrules"] in (HRuleStyle.ALL, HRuleStyle.FRAME):
                 bits.append(self._stringify_hrule(options, "top_"))
                 if options["title"] and options["vrules"] in (
-                    VRuleStyle.ALL,
-                    VRuleStyle.FRAME,
+                        VRuleStyle.ALL,
+                        VRuleStyle.FRAME,
                 ):
                     left_j_len = len(self.left_junction_char)
                     right_j_len = len(self.right_junction_char)

brainstate/util/_struct.py

@@ -21,10 +21,11 @@ from __future__ import annotations
 
 import collections
 import dataclasses
-import jax
 from collections.abc import Hashable, Mapping
 from types import MappingProxyType
 from typing import Any, TypeVar
+
+import jax
 from typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet
 
 __all__ = [

brainstate/util/filter.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import builtins
-
 import dataclasses
 import typing
 from typing import TYPE_CHECKING

{brainstate-0.1.0.post20250413.dist-info → brainstate-0.1.0.post20250422.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250413
+Version: 0.1.0.post20250422
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers
@@ -15,7 +15,6 @@ Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -33,6 +32,7 @@ Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
 Requires-Dist: brainunit (>=0.0.4)
+Requires-Dist: brainevent
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
 Provides-Extra: cuda12
@@ -56,7 +56,7 @@ Requires-Dist: jaxlib[tpu] ; extra == 'tpu'
 <p align="center">
 	<a href="https://pypi.org/project/brainstate/"><img alt="Supported Python Version" src="https://img.shields.io/pypi/pyversions/brainstate"></a>
 	<a href="https://github.com/chaobrain/brainstate/blob/main/LICENSE"><img alt="LICENSE" src="https://img.shields.io/badge/License-Apache%202.0-blue.svg"></a>
-  	<a href='https://brainstate.readthedocs.io/en/latest/?badge=latest'>
+  	<a href='https://brainstate.readthedocs.io/?badge=latest'>
         <img src='https://readthedocs.org/projects/brainstate/badge/?version=latest' alt='Documentation Status' />
     </a>
     <a href="https://badge.fury.io/py/brainstate"><img alt="PyPI version" src="https://badge.fury.io/py/brainstate.svg"></a>