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

brainstate/compile/_make_jaxpr.py

@@ -55,12 +55,10 @@ 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
@@ -68,18 +66,18 @@ 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._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__ = [
@@ -125,8 +123,8 @@ def _new_jax_trace():
     return frame, trace
 
 
-def _init_state_trace_stack() -> StateTraceStack:
-    state_trace: StateTraceStack = StateTraceStack()
+def _init_state_trace_stack(name) -> StateTraceStack:
+    state_trace: StateTraceStack = StateTraceStack(name=name)
 
     if jax.__version_info__ < (0, 4, 36):
         # Should be within the calling of ``jax.make_jaxpr()``
@@ -141,7 +139,7 @@ def _init_state_trace_stack() -> StateTraceStack:
         return state_trace
 
 
-class StatefulFunction(object):
+class StatefulFunction(PrettyObject):
     """
     A wrapper class for a function that collects the states that are read and written by the function. The states are
     collected by the function and returned as a StateDictManager instance. The StateDictManager instance can be used to
@@ -189,6 +187,7 @@ class StatefulFunction(object):
         abstracted_axes: Optional[Any] = None,
         state_returns: Union[str, Tuple[str, ...]] = ('read', 'write'),
         cache_type: Optional[str] = None,
+        name: Optional[str] = None,
     ):
         # explicit parameters
         self.fun = fun
@@ -197,6 +196,7 @@ class StatefulFunction(object):
         self.abstracted_axes = abstracted_axes
         self.state_returns = tuple(state_returns) if isinstance(state_returns, (tuple, list)) else (state_returns,)
         assert cache_type in [None, 'jit']
+        self.name = name
 
         # implicit parameters
         self.cache_type = cache_type
@@ -205,12 +205,10 @@ class StatefulFunction(object):
         self._cached_jaxpr_out_tree: Dict[Any, PyTree] = dict()
         self._cached_state_trace: Dict[Any, StateTraceStack] = dict()
 
-    def __repr__(self) -> str:
-        return (f"{self.__class__.__name__}("
-                f"static_argnums={self.static_argnums}, "
-                f"axis_env={self.axis_env}, "
-                f"abstracted_axes={self.abstracted_axes}, "
-                f"state_returns={self.state_returns})")
+    def __pretty_repr_item__(self, k, v):
+        if k.startswith('_'):
+            return None
+        return k, v
 
     def get_jaxpr(self, cache_key: Hashable = ()) -> jax.core.ClosedJaxpr:
         """
@@ -388,7 +386,7 @@ class StatefulFunction(object):
           A tuple of the states that are read and written by the function and the output of the function.
         """
         # state trace
-        state_trace = _init_state_trace_stack()
+        state_trace = _init_state_trace_stack(self.name)
         self._cached_state_trace[cache_key] = state_trace
         with state_trace:
             out = self.fun(*args, **kwargs)
@@ -497,11 +495,7 @@ class StatefulFunction(object):
         """
         state_trace = self.get_state_trace(self.get_arg_cache_key(*args, **kwargs))
         state_vals, out = self.jaxpr_call([st.value for st in state_trace.states], *args, **kwargs)
-        for st, written, val in zip(state_trace.states, state_trace.been_writen, state_vals):
-            if written:
-                st.value = val
-            else:
-                st.restore_value(val)
+        state_trace.assign_state_vals(state_vals)
         return out
 
 
@@ -592,7 +586,15 @@ def make_jaxpr(
       in (g,) }
     """
 
-    stateful_fun = StatefulFunction(fun, static_argnums, axis_env, abstracted_axes, state_returns)
+    stateful_fun = StatefulFunction(
+        fun,
+        static_argnums=static_argnums,
+        axis_env=axis_env,
+        abstracted_axes=abstracted_axes,
+        state_returns=state_returns,
+        name='make_jaxpr'
+
+    )
 
     @wraps(fun)
     def make_jaxpr_f(*args, **kwargs):

brainstate/compile/_make_jaxpr_test.py

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

brainstate/compile/_progress_bar.py

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

brainstate/compile/_unvmap.py

@@ -19,6 +19,7 @@ import jax.core
 import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
+
 from brainstate._utils import set_module_as
 
 if jax.__version_info__ < (0, 4, 38):

brainstate/compile/_util.py

@@ -132,8 +132,10 @@ def wrap_single_fun(
         assert len(been_writen) == len(writen_state_vals) == len(read_state_vals)
 
         # collect all written and read states
-        state_vals = [written_val if written else read_val
-                      for written, written_val, read_val in zip(been_writen, writen_state_vals, read_state_vals)]
+        state_vals = [
+            written_val if written else read_val
+            for written, written_val, read_val in zip(been_writen, writen_state_vals, read_state_vals)
+        ]
 
         # call the jaxpr
         state_vals, (carry, out) = stateful_fun.jaxpr_call(state_vals, carry, inputs)

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,8 @@
 # ==============================================================================
 
 
-import unittest
-
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 

brainstate/functional/_activations.py

@@ -20,11 +20,10 @@ 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,10 +15,9 @@
 
 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,10 +15,9 @@
 
 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/functional/_spikes.py

@@ -27,53 +27,169 @@ __all__ = [
 
 
 def spike_bitwise_or(x, y):
-    """Bitwise OR operation for spike tensors."""
+    """
+    Perform a bitwise OR operation on spike tensors.
+
+    This function computes the OR operation between two spike tensors.
+    The OR operation is implemented using the formula: x + y - x * y,
+    which is equivalent to the OR operation for binary values.
+
+    Args:
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise OR operation applied to the input tensors.
+               The output tensor has the same shape as the input tensors.
+
+    Note:
+        This operation assumes that the input tensors contain binary (0 or 1) values.
+        For non-binary inputs, the behavior may not correspond to a true bitwise OR.
+    """
     return x + y - x * y
 
 
 def spike_bitwise_and(x, y):
-    """Bitwise AND operation for spike tensors."""
+    """
+    Perform a bitwise AND operation on spike tensors.
+
+    This function computes the AND operation between two spike tensors.
+    The AND operation is equivalent to element-wise multiplication for binary values.
+
+    Args:
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise AND operation applied to the input tensors.
+               The output tensor has the same shape as the input tensors.
+
+    Note:
+        This operation is implemented using element-wise multiplication (x * y),
+        which is equivalent to the AND operation for binary values.
+    """
     return x * y
 
 
 def spike_bitwise_iand(x, y):
-    """Bitwise IAND operation for spike tensors."""
+    """
+    Perform a bitwise IAND (Inverse AND) operation on spike tensors.
+
+    This function computes the Inverse AND (IAND) operation between two spike tensors.
+    IAND is defined as (NOT x) AND y.
+
+    Args:
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise IAND operation applied to the input tensors.
+               The output tensor has the same shape as the input tensors.
+
+    Note:
+        This operation is implemented using the formula: (1 - x) * y,
+        which is equivalent to the IAND operation for binary values.
+    """
     return (1 - x) * y
 
 
 def spike_bitwise_not(x):
-    """Bitwise NOT operation for spike tensors."""
+    """
+    Perform a bitwise NOT operation on spike tensors.
+
+    This function computes the NOT operation on a spike tensor.
+    The NOT operation inverts the binary values in the tensor.
+
+    Args:
+        x (Tensor): The input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise NOT operation applied to the input tensor.
+               The output tensor has the same shape as the input tensor.
+
+    Note:
+        This operation is implemented using the formula: 1 - x,
+        which is equivalent to the NOT operation for binary values.
+    """
     return 1 - x
 
 
 def spike_bitwise_xor(x, y):
-    """Bitwise XOR operation for spike tensors."""
+    """
+    Perform a bitwise XOR operation on spike tensors.
+
+    This function computes the XOR operation between two spike tensors.
+    XOR is defined as (x OR y) AND NOT (x AND y).
+
+    Args:
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise XOR operation applied to the input tensors.
+               The output tensor has the same shape as the input tensors.
+
+    Note:
+        This operation is implemented using the formula: x + y - 2 * x * y,
+        which is equivalent to the XOR operation for binary values.
+    """
     return x + y - 2 * x * y
 
 
 def spike_bitwise_ixor(x, y):
-    """Bitwise IXOR operation for spike tensors."""
+    """
+    Perform a bitwise IXOR (Inverse XOR) operation on spike tensors.
+
+    This function computes the Inverse XOR (IXOR) operation between two spike tensors.
+    IXOR is defined as (x AND NOT y) OR (NOT x AND y).
+
+    Args:
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+
+    Returns:
+        Tensor: The result of the bitwise IXOR operation applied to the input tensors.
+               The output tensor has the same shape as the input tensors.
+
+    Note:
+        This operation is implemented using the formula: x * (1 - y) + (1 - x) * y,
+        which is equivalent to the IXOR operation for binary values.
+    """
     return x * (1 - y) + (1 - x) * y
 
 
 def spike_bitwise(x, y, op: str):
-    r"""Bitwise operation for spike tensors.
-
-    .. math::
+    """
+    Perform bitwise operations on spike tensors.
 
-       \begin{array}{ccc}
-        \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
-        \hline \text { ADD } & x+y & x+y \\
-        \text { AND } & x \cap y & x \cdot y \\
-        \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
-        \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
-        \hline
-        \end{array}
+    This function applies various bitwise operations on spike tensors based on the specified operation.
+    It supports 'or', 'and', 'iand', 'xor', and 'ixor' operations.
 
     Args:
-      x: A spike tensor.
-      y: A spike tensor.
-      op: A string indicating the bitwise operation to perform.
+        x (Tensor): The first input spike tensor.
+        y (Tensor): The second input spike tensor.
+        op (str): A string indicating the bitwise operation to perform.
+            Supported operations are 'or', 'and', 'iand', 'xor', and 'ixor'.
+
+    Returns:
+        Tensor: The result of the bitwise operation applied to the input tensors.
+
+    Raises:
+        NotImplementedError: If an unsupported bitwise operation is specified.
+
+    Note:
+        The function uses the following mathematical expressions for different operations:
+        
+        .. math::
+
+           \begin{array}{ccc}
+            \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
+            \hline \text { ADD } & x+y & x+y \\
+            \text { AND } & x \cap y & x \cdot y \\
+            \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
+            \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
+            \hline
+            \end{array}
     """
     if op == 'or':
         return spike_bitwise_or(x, y)

brainstate/graph/_graph_node.py

@@ -27,7 +27,7 @@ import numpy as np
 
 from brainstate._state import State, TreefyState
 from brainstate.typing import Key
-from brainstate.util._pretty_repr import PrettyRepr, yield_unique_pretty_repr_items, PrettyType, PrettyAttr
+from brainstate.util._pretty_pytree import PrettyObject
 from ._graph_operation import register_graph_node_type
 
 __all__ = [
@@ -46,7 +46,7 @@ class GraphNodeMeta(ABCMeta):
             return node
 
 
-class Node(PrettyRepr, metaclass=GraphNodeMeta):
+class Node(PrettyObject, metaclass=GraphNodeMeta):
     """
     Base class for all graph nodes.
 
@@ -84,47 +84,6 @@ class Node(PrettyRepr, metaclass=GraphNodeMeta):
         state = deepcopy(state)
         return treefy_merge(graphdef, state)
 
-    def __pretty_repr__(self):
-        """
-        Pretty repr for the object.
-        """
-        yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
-
-    def __treescope_repr__(self, path, subtree_renderer):
-        """
-        Treescope repr for the object.
-        """
-        children = {}
-        for name, value in vars(self).items():
-            name, value = self.__leaf_fn__(name, value)
-            if name.startswith('_'):
-                continue
-            children[name] = value
-        import treescope  # type: ignore[import-not-found,import-untyped]
-        return treescope.repr_lib.render_object_constructor(
-            object_type=type(self),
-            attributes=children,
-            path=path,
-            subtree_renderer=subtree_renderer,
-            color=treescope.formatting_util.color_from_string(type(self).__qualname__)
-        )
-
-    def __leaf_fn__(self, leaf, value):
-        return leaf, value
-
-
-def _default_repr_object(node: Node):
-    yield PrettyType(type=type(node))
-
-
-def _default_repr_attr(node: Node):
-    for name, value in vars(node).items():
-        name, value = node.__leaf_fn__(name, value)
-        if name.startswith('_'):
-            continue
-        # value = jax.tree.map(_to_shape_dtype, value, is_leaf=lambda x: isinstance(x, u.Quantity))
-        yield PrettyAttr(name, repr(value))
-
 
 class String:
     def __init__(self, msg):

brainstate/graph/_graph_operation.py

@@ -18,21 +18,20 @@
 from __future__ import annotations
 
 import dataclasses
-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 import (Any, Callable, Generic, Iterable, Iterator, Mapping, MutableMapping,
+                    Sequence, Type, TypeVar, Union, Hashable, Tuple, Dict, Optional, overload)
 from typing_extensions import TypeGuard, Unpack
 
 from brainstate._state import State, TreefyState
 from brainstate._utils import set_module_as
 from brainstate.typing import PathParts, Filter, Predicate, Key
 from brainstate.util._caller import ApplyCaller, CallableProxy, DelayedAccessor
-from brainstate.util._dict import NestedDict, FlattedDict, PrettyDict
-from brainstate.util._filter import to_predicate
+from brainstate.util._pretty_pytree import NestedDict, FlattedDict, PrettyDict
 from brainstate.util._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, PrettyMapping, MappingReprMixin
 from brainstate.util._struct import FrozenDict
+from brainstate.util.filter import to_predicate
 
 _max_int = np.iinfo(np.int32).max
 
@@ -347,21 +346,6 @@ class NodeDef(GraphDef[Node], PrettyRepr):
         yield PrettyAttr('metadata', self.metadata)
         yield PrettyAttr('index_mapping', PrettyMapping(self.index_mapping) if self.index_mapping is not None else None)
 
-    def __treescope_repr__(self, path, subtree_renderer):
-        import treescope  # type: ignore[import-not-found,import-untyped]
-        return treescope.repr_lib.render_object_constructor(
-            object_type=type(self),
-            attributes={'type': self.type,
-                        'index': self.index,
-                        'attributes': self.attributes,
-                        'subgraphs': dict(self.subgraphs),
-                        'static_fields': dict(self.static_fields),
-                        'leaves': dict(self.leaves),
-                        'metadata': self.metadata, },
-            path=path,
-            subtree_renderer=subtree_renderer,
-        )
-
     def apply(
         self,
         state_map: GraphStateMapping,

brainstate/graph/_graph_operation_test.py

@@ -15,13 +15,12 @@
 
 from __future__ import annotations
 
-import unittest
-from collections.abc import Callable
-from threading import Thread
-
 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 brainstate as bst
 

brainstate/init/_base.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-from typing import Optional, Tuple
-
 import numpy as np
+from typing import Optional, Tuple
 
 from brainstate.util import PrettyRepr, PrettyType, PrettyAttr
 

brainstate/init/_generic.py

@@ -16,11 +16,10 @@
 # -*- 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

@@ -17,6 +17,8 @@
 from . import metrics
 from ._collective_ops import *
 from ._collective_ops import __all__ as collective_ops_all
+from ._common import *
+from ._common import __all__ as common_all
 from ._dyn_impl import *
 from ._dyn_impl import __all__ as dyn_impl_all
 from ._dynamics import *
@@ -33,6 +35,7 @@ from ._module import __all__ as module_all
 __all__ = (
     ['metrics']
     + collective_ops_all
+    + common_all
     + dyn_impl_all
     + dynamics_all
     + elementwise_all
@@ -43,6 +46,7 @@ __all__ = (
 
 del (
     collective_ops_all,
+    common_all,
     dyn_impl_all,
     dynamics_all,
     elementwise_all,