brainstate 0.1.0.post20250212__py2.py3-none-any.whl → 0.1.0.post20250217__py2.py3-none-any.whl

brainstate/random/_rand_seed.py

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

brainstate/random/_rand_seed_test.py

@@ -15,10 +15,9 @@
 
 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,17 +16,16 @@
 # -*- 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,7 +21,7 @@ import jax.numpy as jnp
 import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
-from brainstate.util._dict import PrettyObject
+from brainstate.util._pretty_pytree import PrettyObject
 
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import Primitive
@@ -79,7 +79,10 @@ def _heaviside_imp(x, dx):
 
 
 def _heaviside_batching(args, axes):
-    return heaviside_p.bind(*args), tuple(axes)
+    x, dx = args
+    if axes[0] != axes[1]:
+        dx = batching.moveaxis(dx, axes[1], axes[0])
+    return heaviside_p.bind(x, dx), tuple([axes[0]])
 
 
 def _heaviside_jvp(primals, tangents):
@@ -129,9 +132,6 @@ class Surrogate(PrettyObject):
         dx = self.surrogate_grad(x)
         return heaviside_p.bind(x, dx)[0]
 
-    def __repr__(self):
-        return f'{self.__class__.__name__}()'
-
     def surrogate_fun(self, x) -> jax.Array:
         """The surrogate function."""
         raise NotImplementedError

brainstate/transform.py

@@ -15,6 +15,3 @@
 
 # alias for compilation and augmentation functions
 
-from .compile import *
-from .augment import *
-

brainstate/typing.py

@@ -16,13 +16,17 @@
 from __future__ import annotations
 
 import builtins
+
+import brainunit as u
 import functools as ft
 import importlib
 import inspect
-
-import brainunit as u
 import jax
 import numpy as np
+from typing import (
+    Any, Callable, Hashable, List, Protocol, Tuple, TypeVar, Union,
+    runtime_checkable, TYPE_CHECKING, Generic, Sequence
+)
 
 tp = importlib.import_module("typing")
 
@@ -41,35 +45,35 @@ __all__ = [
     'Missing',
 ]
 
-K = tp.TypeVar('K')
+K = TypeVar('K')
 
 
-@tp.runtime_checkable
-class Key(tp.Hashable, tp.Protocol):
+@runtime_checkable
+class Key(Hashable, Protocol):
     def __lt__(self: K, value: K, /) -> bool:
         ...
 
 
-Ellipsis = builtins.ellipsis if tp.TYPE_CHECKING else tp.Any
+Ellipsis = builtins.ellipsis if TYPE_CHECKING else Any
 
-PathParts = tp.Tuple[Key, ...]
-Predicate = tp.Callable[[PathParts, tp.Any], bool]
-FilterLiteral = tp.Union[type, str, Predicate, bool, Ellipsis, None]
-Filter = tp.Union[FilterLiteral, tp.Tuple['Filter', ...], tp.List['Filter']]
+PathParts = Tuple[Key, ...]
+Predicate = Callable[[PathParts, Any], bool]
+FilterLiteral = Union[type, str, Predicate, bool, Ellipsis, None]
+Filter = Union[FilterLiteral, Tuple['Filter', ...], List['Filter']]
 
-_T = tp.TypeVar("_T")
+_T = TypeVar("_T")
 
-_Annotation = tp.TypeVar("_Annotation")
+_Annotation = TypeVar("_Annotation")
 
 
-class _Array(tp.Generic[_Annotation]):
+class _Array(Generic[_Annotation]):
     pass
 
 
 _Array.__module__ = "builtins"
 
 
-def _item_to_str(item: tp.Union[str, type, slice]) -> str:
+def _item_to_str(item: Union[str, type, slice]) -> str:
     if isinstance(item, slice):
         if item.step is not None:
             raise NotImplementedError
@@ -83,7 +87,7 @@ def _item_to_str(item: tp.Union[str, type, slice]) -> str:
 
 
 def _maybe_tuple_to_str(
-    item: tp.Union[str, type, slice, tp.Tuple[tp.Union[str, type, slice], ...]]
+    item: Union[str, type, slice, Tuple[Union[str, type, slice], ...]]
 ) -> str:
     if isinstance(item, tuple):
         if len(item) == 0:
@@ -113,7 +117,7 @@ class Array:
 Array.__module__ = "builtins"
 
 
-class _FakePyTree(tp.Generic[_T]):
+class _FakePyTree(Generic[_T]):
     pass
 
 
@@ -255,11 +259,10 @@ f. A structure can end with a `...`, to denote that the PyTree must be a prefix
     cases, all named pieces must already have been seen and their structures bound.
 """  # noqa: E501
 
-Size = tp.Union[int, tp.Sequence[int]]
-Axes = tp.Union[int, tp.Sequence[int]]
-SeedOrKey = tp.Union[int, jax.Array, np.ndarray]
-Shape = tp.Sequence[int]
-
+Size = Union[int, Sequence[int]]
+Axes = Union[int, Sequence[int]]
+SeedOrKey = Union[int, jax.Array, np.ndarray]
+Shape = Sequence[int]
 
 # --- Array --- #
 
@@ -267,7 +270,7 @@ Shape = tp.Sequence[int]
 # standard JAX array (i.e. not including future non-standard array types like
 # KeyArray and BInt). It's different than np.typing.ArrayLike in that it doesn't
 # accept arbitrary sequences, nor does it accept string data.
-ArrayLike = tp.Union[
+ArrayLike = Union[
     jax.Array,  # JAX array type
     np.ndarray,  # NumPy array type
     np.bool_, np.number,  # NumPy scalar types
@@ -281,7 +284,7 @@ ArrayLike = tp.Union[
 DType = np.dtype
 
 
-class SupportsDType(tp.Protocol):
+class SupportsDType(Protocol):
     @property
     def dtype(self) -> DType: ...
 
@@ -291,9 +294,9 @@ class SupportsDType(tp.Protocol):
 # because JAX doesn't support objects or structured dtypes.
 # Unlike np.typing.DTypeLike, we exclude None, and instead require
 # explicit annotations when None is acceptable.
-DTypeLike = tp.Union[
+DTypeLike = Union[
     str,  # like 'float32', 'int32'
-    type[tp.Any],  # like np.float32, np.int32, float, int
+    type[Any],  # like np.float32, np.int32, float, int
     np.dtype,  # like np.dtype('float32'), np.dtype('int32')
     SupportsDType,  # like jnp.float32, jnp.int32
 ]

brainstate/util/__init__.py

@@ -13,36 +13,38 @@
 # limitations under the License.
 # ==============================================================================
 
-from ._dict import *
-from ._dict import __all__ as _mapping_all
+from . import filter
 from ._error import *
 from ._error import __all__ as _error_all
-from ._filter import *
-from ._filter import __all__ as _filter_all
 from ._others import *
 from ._others import __all__ as _others_all
+from ._pretty_pytree import *
+from ._pretty_pytree import __all__ as _mapping_all
 from ._pretty_repr import *
 from ._pretty_repr import __all__ as _pretty_repr_all
+from ._pretty_table import *
+from ._pretty_table import __all__ as _table_all
 from ._scaling import *
 from ._scaling import __all__ as _mem_scale_all
 from ._struct import *
 from ._struct import __all__ as _struct_all
 
 __all__ = (
-    _others_all
+    ['filter']
+    + _others_all
     + _mem_scale_all
-    + _filter_all
     + _pretty_repr_all
     + _struct_all
     + _error_all
     + _mapping_all
+    + _table_all
 )
 del (
     _others_all,
     _mem_scale_all,
-    _filter_all,
     _pretty_repr_all,
     _struct_all,
     _error_all,
     _mapping_all,
+    _table_all,
 )

brainstate/util/_caller.py

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

brainstate/util/_error.py

@@ -21,8 +21,35 @@ __all__ = [
 
 
 class BrainStateError(Exception):
+    """
+    A custom exception class for BrainState-related errors.
+
+    This exception is raised when a BrainState-specific error occurs during
+    the execution of the program. It serves as a base class for more specific
+    BrainState exceptions.
+
+    Attributes:
+        Inherits all attributes from the built-in Exception class.
+
+    Usage::
+
+        raise BrainStateError("A BrainState-specific error occurred.")
+    """
     pass
 
 
 class TraceContextError(BrainStateError):
+    """
+    A custom exception class for trace context-related errors in BrainState.
+
+    This exception is raised when an error occurs specifically related to
+    trace context operations or manipulations within the BrainState framework.
+
+    Attributes:
+        Inherits all attributes from the BrainStateError class.
+
+    Usage::
+
+        raise TraceContextError("An error occurred while handling trace context.")
+    """
     pass

brainstate/util/_others.py

@@ -15,20 +15,21 @@
 
 from __future__ import annotations
 
+import gc
+
 import copy
 import functools
-import gc
+import jax
 import threading
 import types
 from collections.abc import Iterable
-from typing import Any, Callable, Tuple, Union, Dict
-
-import jax
 from jax.lib import xla_bridge
+from typing import Any, Callable, Tuple, Union, Dict
 
 from brainstate._utils import set_module_as
 
 __all__ = [
+    'split_total',
     'clear_buffer_memory',
     'not_instance_eval',
     'is_instance_eval',
@@ -37,6 +38,61 @@ __all__ = [
 ]
 
 
+def split_total(
+    total: int,
+    fraction: Union[int, float],
+) -> int:
+    """
+   Calculate the number of epochs for simulation based on a total and a fraction.
+
+   This function determines the number of epochs to simulate given a total number
+   of epochs and either a fraction or a specific number of epochs to run.
+
+   Parameters:
+   -----------
+   total : int
+       The total number of epochs. Must be a positive integer.
+   fraction : Union[int, float]
+       If ``float``: A value between 0 and 1 representing the fraction of total epochs to run.
+       If ``int``: The specific number of epochs to run, must not exceed the total.
+
+   Returns:
+   --------
+   int
+       The calculated number of epochs to simulate.
+
+   Raises:
+   -------
+   ValueError
+       If total is not positive, fraction is negative, or if fraction as float is > 1
+       or as int is > total.
+   AssertionError
+       If total is not an integer.
+   """
+    assert isinstance(total, int), "Length must be an integer."
+    if total <= 0:
+        raise ValueError("'total' must be a positive integer.")
+    if fraction < 0:
+        raise ValueError("'fraction' value cannot be negative.")
+
+    if isinstance(fraction, float):
+        if fraction < 0:
+            raise ValueError("'fraction' value cannot be negative.")
+        if fraction > 1:
+            raise ValueError("'fraction' value cannot be greater than 1.")
+        return int(total * fraction)
+
+    elif isinstance(fraction, int):
+        if fraction < 0:
+            raise ValueError("'fraction' value cannot be negative.")
+        if fraction > total:
+            raise ValueError("'fraction' value cannot be greater than total.")
+        return fraction
+
+    else:
+        raise ValueError("'fraction' must be an integer or float.")
+
+
 class NameContext(threading.local):
     def __init__(self):
         self.typed_names: Dict[str, int] = {}
@@ -249,17 +305,6 @@ class DictManager(dict):
         else:
             raise ValueError(f'Unsupported method: {by}')
 
-    def union_by_value_ids(self, other: dict):
-        """
-        Union the stack by the value ids.
-
-        Args:
-          other:
-
-        Returns:
-
-        """
-
     def __add__(self, other: dict):
         """
         Compose other instance of dict.

brainstate/util/{_dict.py → _pretty_pytree.py}

@@ -18,14 +18,14 @@
 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 ._filter import to_predicate
 from ._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, yield_unique_pretty_repr_items, pretty_repr_object
 from ._struct import dataclass
+from .filter import to_predicate
 
 __all__ = [
     'PrettyDict',

brainstate/util/{_dict_test.py → _pretty_pytree_test.py}

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

brainstate/util/_pretty_repr.py

@@ -21,7 +21,7 @@ import dataclasses
 import threading
 from abc import ABC, abstractmethod
 from functools import partial
-from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional, Sequence
+from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional
 
 __all__ = [
     'yield_unique_pretty_repr_items',
@@ -328,4 +328,3 @@ def yield_unique_pretty_repr_items(
     finally:
         if clear_seen:
             CONTEXT.seen_modules_repr = None
-