brainstate 0.1.0.post20250503__py2.py3-none-any.whl → 0.1.1__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
 
 from . import augment
 from . import compile

brainstate/_compatible_import.py

@@ -35,8 +35,8 @@ __all__ = [
     'safe_map',
     'safe_zip',
     'unzip2',
-    'unzip3',
     'wraps',
+    'Device',
 ]
 
 T = TypeVar("T")
@@ -48,6 +48,11 @@ brainevent_installed = importlib.util.find_spec('brainevent') is not None
 
 from jax.core import get_aval, Tracer
 
+if jax.__version_info__ < (0, 5, 0):
+    from jax.lib.xla_client import Device
+else:
+    from jax import Device
+
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import ClosedJaxpr, extend_axis_env_nd, Primitive, jaxpr_as_fun
 else:
@@ -84,8 +89,7 @@ else:
         return list(zip(*args))
 
 
-    def unzip2(xys: Iterable[tuple[T1, T2]]
-               ) -> tuple[tuple[T1, ...], tuple[T2, ...]]:
+    def unzip2(xys: Iterable[tuple[T1, T2]]) -> tuple[tuple[T1, ...], tuple[T2, ...]]:
         """Unzip sequence of length-2 tuples into two tuples."""
         # Note: we deliberately don't use zip(*xys) because it is lazily evaluated,
         # is too permissive about inputs, and does not guarantee a length-2 output.

brainstate/_state.py

@@ -133,183 +133,6 @@ def check_state_value_tree(val: bool = True) -> Generator[None, None, None]:
         TRACE_CONTEXT.tree_check.pop()
 
 
-class StateCatcher(PrettyObject):
-    """
-    The catcher to catch and manage new states.
-
-    This class provides functionality to collect and tag new State objects.
-    It ensures that each state is only added once and assigns a tag to each state.
-
-    Attributes:
-        state_tag (str): A string identifier used to tag the caught states.
-        state_ids (set): A set of state IDs to ensure uniqueness.
-        states (list): A list to store the caught State objects.
-    """
-
-    def __init__(
-        self,
-        state_tag: str,
-        state_to_exclude: Filter = Nothing()
-    ):
-        """
-        Initialize a new Catcher instance.
-
-        Args:
-            state_tag (str): The tag to be assigned to caught states.
-            state_to_exclude (Filter, optional): A filter to exclude states from being caught.
-        """
-        if state_to_exclude is None:
-            state_to_exclude = Nothing()
-        self.state_to_exclude = state_to_exclude
-        self.state_tag = state_tag
-        self.state_ids = set()
-        self.states = []
-
-    def get_state_values(self) -> List[PyTree]:
-        """
-        Get the values of the caught states.
-
-        Returns:
-            list: A list of values of the caught states.
-        """
-        return [state.value for state in self.states]
-
-    def get_states(self) -> List[State]:
-        """
-        Get the caught states.
-
-        Returns:
-            list: A list of the caught states.
-        """
-        return self.states
-
-    def append(self, state: State):
-        """
-        Add a new state to the catcher if it hasn't been added before.
-
-        This method adds the state to the internal list, records its ID,
-        and assigns the catcher's tag to the state.
-
-        Args:
-            state (State): The State object to be added.
-        """
-        if self.state_to_exclude((), state):
-            return
-        if id(state) not in self.state_ids:
-            self.state_ids.add(id(state))
-            self.states.append(state)
-            state.tag = self.state_tag
-
-    def __iter__(self):
-        """
-        Allow iteration over the caught states.
-
-        Returns:
-            iterator: An iterator over the list of caught states.
-        """
-        return iter(self.states)
-
-    def __len__(self):
-        """
-        Return the number of caught states.
-
-        Returns:
-            int: The number of caught states.
-        """
-        return len(self.states)
-
-    def __getitem__(self, index):
-        """
-        Get a state by index.
-
-        Args:
-            index (int): The index of the state to retrieve.
-
-        Returns:
-            State: The state at the specified index.
-        """
-        return self.states[index]
-
-    def clear(self):
-        """
-        Clear all caught states.
-        """
-        self.state_ids.clear()
-        self.states.clear()
-
-    def get_by_tag(self, tag: str):
-        """
-        Get all states with a specific tag.
-
-        Args:
-            tag (str): The tag to filter by.
-
-        Returns:
-            list: A list of states with the specified tag.
-        """
-        return [state for state in self.states if state.tag == tag]
-
-    def remove(self, state: State):
-        """
-        Remove a specific state from the catcher.
-
-        Args:
-            state (State): The state to remove.
-        """
-        if id(state) in self.state_ids:
-            self.state_ids.remove(id(state))
-            self.states.remove(state)
-
-    def __contains__(self, state: State):
-        """
-        Check if a state is in the catcher.
-
-        Args:
-            state (State): The state to check for.
-
-        Returns:
-            bool: True if the state is in the catcher, False otherwise.
-        """
-        return id(state) in self.state_ids
-
-
-@contextlib.contextmanager
-def catch_new_states(
-    state_tag: str = None,
-    state_to_exclude: Filter = Nothing()
-) -> Generator[StateCatcher, None, None]:
-    """
-    A context manager that catches and tracks new states created within its scope.
-
-    This function creates a new Catcher object and adds it to the TRACE_CONTEXT's
-    new_state_catcher list. It allows for tracking and managing new states created
-    within the context.
-
-    Args:
-        state_tag (str, optional): A string tag to associate with the caught states.
-            Defaults to None.
-        state_to_exclude (Filter, optional): A filter object to specify which states
-            should be excluded from catching. Defaults to Nothing(), which excludes no states.
-
-    Yields:
-        Catcher: A Catcher object that can be used to access and manage the
-            newly created states within the context.
-
-    Example::
-    
-        with catch_new_states("my_tag") as catcher:
-            # Create new states here
-            # They will be caught and tagged with "my_tag"
-        # Access caught states through catcher object
-    """
-    try:
-        catcher = StateCatcher(state_tag=state_tag, state_to_exclude=state_to_exclude)
-        TRACE_CONTEXT.new_state_catcher.append(catcher)
-        yield catcher
-    finally:
-        TRACE_CONTEXT.new_state_catcher.pop()
-
-
 def maybe_state(val: Any) -> Any:
     """
     Extracts the value from a State object if given, otherwise returns the input value.
@@ -1660,3 +1483,180 @@ jax.tree_util.register_pytree_with_keys(
     _state_ref_unflatten,  # type: ignore
     flatten_func=partial(_state_ref_flatten, with_keys=False),  # type: ignore
 )
+
+
+class StateCatcher(PrettyObject):
+    """
+    The catcher to catch and manage new states.
+
+    This class provides functionality to collect and tag new State objects.
+    It ensures that each state is only added once and assigns a tag to each state.
+
+    Attributes:
+        state_tag (str): A string identifier used to tag the caught states.
+        state_ids (set): A set of state IDs to ensure uniqueness.
+        states (list): A list to store the caught State objects.
+    """
+
+    def __init__(
+        self,
+        state_tag: str,
+        state_to_exclude: Filter = Nothing()
+    ):
+        """
+        Initialize a new Catcher instance.
+
+        Args:
+            state_tag (str): The tag to be assigned to caught states.
+            state_to_exclude (Filter, optional): A filter to exclude states from being caught.
+        """
+        if state_to_exclude is None:
+            state_to_exclude = Nothing()
+        self.state_to_exclude = state_to_exclude
+        self.state_tag = state_tag
+        self.state_ids = set()
+        self.states = []
+
+    def get_state_values(self) -> List[PyTree]:
+        """
+        Get the values of the caught states.
+
+        Returns:
+            list: A list of values of the caught states.
+        """
+        return [state.value for state in self.states]
+
+    def get_states(self) -> List[State]:
+        """
+        Get the caught states.
+
+        Returns:
+            list: A list of the caught states.
+        """
+        return self.states
+
+    def append(self, state: State):
+        """
+        Add a new state to the catcher if it hasn't been added before.
+
+        This method adds the state to the internal list, records its ID,
+        and assigns the catcher's tag to the state.
+
+        Args:
+            state (State): The State object to be added.
+        """
+        if self.state_to_exclude((), state):
+            return
+        if id(state) not in self.state_ids:
+            self.state_ids.add(id(state))
+            self.states.append(state)
+            state.tag = self.state_tag
+
+    def __iter__(self):
+        """
+        Allow iteration over the caught states.
+
+        Returns:
+            iterator: An iterator over the list of caught states.
+        """
+        return iter(self.states)
+
+    def __len__(self):
+        """
+        Return the number of caught states.
+
+        Returns:
+            int: The number of caught states.
+        """
+        return len(self.states)
+
+    def __getitem__(self, index):
+        """
+        Get a state by index.
+
+        Args:
+            index (int): The index of the state to retrieve.
+
+        Returns:
+            State: The state at the specified index.
+        """
+        return self.states[index]
+
+    def clear(self):
+        """
+        Clear all caught states.
+        """
+        self.state_ids.clear()
+        self.states.clear()
+
+    def get_by_tag(self, tag: str):
+        """
+        Get all states with a specific tag.
+
+        Args:
+            tag (str): The tag to filter by.
+
+        Returns:
+            list: A list of states with the specified tag.
+        """
+        return [state for state in self.states if state.tag == tag]
+
+    def remove(self, state: State):
+        """
+        Remove a specific state from the catcher.
+
+        Args:
+            state (State): The state to remove.
+        """
+        if id(state) in self.state_ids:
+            self.state_ids.remove(id(state))
+            self.states.remove(state)
+
+    def __contains__(self, state: State):
+        """
+        Check if a state is in the catcher.
+
+        Args:
+            state (State): The state to check for.
+
+        Returns:
+            bool: True if the state is in the catcher, False otherwise.
+        """
+        return id(state) in self.state_ids
+
+
+@contextlib.contextmanager
+def catch_new_states(
+    state_tag: str = None,
+    state_to_exclude: Filter = Nothing()
+) -> Generator[StateCatcher, None, None]:
+    """
+    A context manager that catches and tracks new states created within its scope.
+
+    This function creates a new Catcher object and adds it to the TRACE_CONTEXT's
+    new_state_catcher list. It allows for tracking and managing new states created
+    within the context.
+
+    Args:
+        state_tag (str, optional): A string tag to associate with the caught states.
+            Defaults to None.
+        state_to_exclude (Filter, optional): A filter object to specify which states
+            should be excluded from catching. Defaults to Nothing(), which excludes no states.
+
+    Yields:
+        Catcher: A Catcher object that can be used to access and manage the
+            newly created states within the context.
+
+    Example::
+
+        with catch_new_states("my_tag") as catcher:
+            # Create new states here
+            # They will be caught and tagged with "my_tag"
+        # Access caught states through catcher object
+    """
+    try:
+        catcher = StateCatcher(state_tag=state_tag, state_to_exclude=state_to_exclude)
+        TRACE_CONTEXT.new_state_catcher.append(catcher)
+        yield catcher
+    finally:
+        TRACE_CONTEXT.new_state_catcher.pop()

brainstate/_utils.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import warnings
 

brainstate/augment/_autograd.py

@@ -27,8 +27,6 @@ The wrapped gradient transformations here are made possible by using the followi
 
 """
 
-from __future__ import annotations
-
 from functools import wraps, partial
 from typing import Union, Callable, Dict, Sequence, Optional, Any, Tuple, TypeVar, Iterator
 

brainstate/augment/_eval_shape.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from typing import Any, TypeVar, Callable, Sequence, Union
 

brainstate/augment/_mapping.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from typing import (
     Any,
@@ -33,6 +31,7 @@ from typing import (
 import jax
 from jax.interpreters.batching import BatchTracer
 
+from brainstate._compatible_import import Device
 from brainstate._state import State, catch_new_states
 from brainstate.compile import scan, StatefulFunction
 from brainstate.random import RandomState, DEFAULT
@@ -700,7 +699,7 @@ def pmap(
     in_axes: Any = 0,
     out_axes: Any = 0,
     static_broadcasted_argnums: int | Iterable[int] = (),
-    devices: Optional[Sequence[jax.Device]] = None,  # noqa: F811
+    devices: Optional[Sequence[Device]] = None,  # noqa: F811
     backend: Optional[str] = None,
     axis_size: Optional[int] = None,
     donate_argnums: int | Iterable[int] = (),

brainstate/augment/_random.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from typing import Callable, Sequence, Union
 

brainstate/compile/_ad_checkpoint.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from typing import Callable, Tuple, Union
 

brainstate/compile/_conditions.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from collections.abc import Callable, Sequence
 
 import jax

brainstate/compile/_error_if.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from functools import partial
 from typing import Callable, Union

brainstate/compile/_jit.py

@@ -13,16 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 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 brainstate._compatible_import import Device
 from brainstate._utils import set_module_as
 from brainstate.typing import Missing
 from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
@@ -94,8 +92,8 @@ def _get_jitted_fun(
         read_state_vals = state_trace.get_read_state_values(True)
 
         # call the jitted function
-        # print('Running ...')
         write_state_vals, outs = jit_fun(state_trace.get_state_values(), *args, **params)
+
         # write the state values back to the states
         write_back_state_values(state_trace, read_state_vals, write_state_vals)
         return outs
@@ -106,8 +104,11 @@ def _get_jitted_fun(
         """
         # clear the cache of the stateful function
         fun.clear_cache()
-        # clear the cache of the jitted function
-        jit_fun.clear_cache()
+        try:
+            # clear the cache of the jitted function
+            jit_fun.clear_cache()
+        except AttributeError:
+            pass
 
     def eval_shape():
         raise NotImplementedError
@@ -165,7 +166,7 @@ def _get_jitted_fun(
     # compile the jitted function
     jitted_fun.compile = compile
 
-    # trace the jitted
+    # trace the jitted function
     jitted_fun.trace = trace
 
     return jitted_fun
@@ -180,7 +181,7 @@ def jit(
     donate_argnums: int | Sequence[int] | None = None,
     donate_argnames: str | Iterable[str] | None = None,
     keep_unused: bool = False,
-    device: xc.Device | None = None,
+    device: Device | None = None,
     backend: str | None = None,
     inline: bool = False,
     abstracted_axes: Any | None = None,

brainstate/compile/_loop_collect_return.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import math
 from functools import wraps
 from typing import Callable, Optional, TypeVar, Tuple, Any

brainstate/compile/_loop_no_collection.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import math
 from typing import Any, Callable, TypeVar
 

brainstate/compile/_make_jaxpr.py

@@ -51,8 +51,6 @@ function.
 
 """
 
-from __future__ import annotations
-
 import functools
 import inspect
 import operator
@@ -64,7 +62,7 @@ import jax
 from jax._src import source_info_util
 from jax._src.linear_util import annotate
 from jax._src.traceback_util import api_boundary
-from jax.api_util import shaped_abstractify
+from jax.api_util import shaped_abstractify, debug_info
 from jax.extend.linear_util import transformation_with_aux, wrap_init
 from jax.interpreters import partial_eval as pe
 
@@ -745,7 +743,7 @@ def _make_jaxpr(
     @wraps(fun)
     @api_boundary
     def make_jaxpr_f(*args, **kwargs):
-        f = wrap_init(fun)
+        f = wrap_init(fun, debug_info=debug_info('make_jaxpr', fun, args, kwargs))
         if static_argnums:
             dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
             f, args = jax.api_util.argnums_partial(f, dyn_argnums, args)
@@ -754,12 +752,12 @@ def _make_jaxpr(
         f, out_tree = _flatten_fun(f, in_tree)
         f = annotate(f, in_type)
         if jax.__version_info__ < (0, 5, 0):
-            debug_info = pe.debug_info(fun, in_tree, out_tree, True, '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(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)
+                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 = ClosedJaxpr(jaxpr, consts)

brainstate/compile/_progress_bar.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import copy
 import importlib.util

brainstate/compile/_unvmap.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 import jax
 import jax.core

brainstate/compile/_util.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from functools import wraps
 from typing import Sequence, Tuple
 

brainstate/environ.py

@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import contextlib
 import dataclasses
 import functools

brainstate/functional/_activations.py

@@ -18,8 +18,6 @@
 Shared neural network activations and other functions.
 """
 
-from __future__ import annotations
-
 from typing import Any, Union, Sequence
 
 import brainunit as u

brainstate/functional/_normalization.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from typing import Optional, Union
 
 import brainunit as u

brainstate/functional/_others.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from functools import partial
 
 import jax