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

brainstate/util/_pretty_pytree.py

@@ -15,8 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from collections import abc
 from typing import TypeVar, Hashable, Union, Iterable, Any, Optional, Tuple, Dict
 
@@ -46,66 +44,6 @@ ExtractValueFn = abc.Callable[[Any], Any]
 SetValueFn = abc.Callable[[V, Any], V]
 
 
-def _repr_object_general(node: PrettyDict):
-    """
-    Generate a general representation of a PrettyDict object.
-
-    This function is used to create a pretty representation of a PrettyDict
-    object, which includes the type of the object and its value separator.
-
-    Args:
-        node (PrettyDict): The PrettyDict object to be represented.
-
-    Yields:
-        PrettyType: A PrettyType object representing the type of the node,
-        with specified value separator, start, and end characters.
-    """
-    yield PrettyType(type(node), value_sep='=', start='(', end=')')
-
-
-def _repr_attribute_general(node):
-    """
-    Generate a pretty representation of the attributes of a node.
-
-    This function iterates over the attributes of a given node and attempts
-    to generate a pretty representation for each attribute. It handles
-    conversion of lists and dictionaries to their pretty representation
-    counterparts and yields a PrettyAttr object for each attribute.
-
-    Args:
-        node: The object whose attributes are to be represented.
-
-    Yields:
-        PrettyAttr: A PrettyAttr object representing the key and value of
-        each attribute in a pretty format.
-    """
-    for k, v in vars(node).items():
-        try:
-            res = node.__pretty_repr_item__(k, v)
-            if res is None:
-                continue
-            k, v = res
-        except AttributeError:
-            pass
-
-        if k is None:
-            continue
-
-        # convert list to PrettyList
-        if isinstance(v, list):
-            v = PrettyList(v)
-
-        # convert dict to PrettyDict
-        if isinstance(v, dict):
-            v = PrettyDict(v)
-
-        # convert PrettyDict to NestedStateRepr
-        if isinstance(v, PrettyDict):
-            v = NestedStateRepr(v)
-
-        yield PrettyAttr(k, v)
-
-
 class PrettyObject(PrettyRepr):
     """
     A class for generating a pretty representation of a tree-like structure.
@@ -284,37 +222,15 @@ def _default_process(x):
     return id(x)
 
 
-class NestedStateRepr(PrettyRepr):
-    def __init__(self, state: PrettyDict):
-        self.state = state
-
-    def __pretty_repr__(self):
-        yield PrettyType('', value_sep=': ', start='{', end='}')
-
-        for r in self.state.__pretty_repr__():
-            if isinstance(r, PrettyType):
-                continue
-            yield r
-
-    def __treescope_repr__(self, path, subtree_renderer):
-        children = {}
-        for k, v in self.state.items():
-            if isinstance(v, PrettyDict):
-                v = NestedStateRepr(v)
-            children[k] = v
-        # Render as the dictionary itself at the same path.
-        return subtree_renderer(children, path=path)
-
-
 class PrettyDict(dict, PrettyRepr):
     __module__ = 'brainstate.util'
 
-    def __getattr__(self, key: K) -> NestedMapping | V:  # type: ignore[misc]
+    def __getattr__(self, key: K):  # type: ignore[misc]
         return self[key]
 
     def treefy_state(self):
         """
-        Convert the ``State`` objects to a reference tree of the state.
+        Convert the :class:`State` objects to a reference tree of the state.
         """
         from brainstate._state import State
         leaves, treedef = jax.tree.flatten(self)
@@ -323,7 +239,7 @@ class PrettyDict(dict, PrettyRepr):
 
     def to_dict(self) -> Dict[K, Dict[K, Any] | V]:
         """
-        Convert the ``PrettyDict`` to a dictionary.
+        Convert the :class:`PrettyDict` to a dictionary.
 
         Returns:
           The dictionary.
@@ -337,24 +253,46 @@ class PrettyDict(dict, PrettyRepr):
     def __pretty_repr__(self):
         yield from yield_unique_pretty_repr_items(self, _default_repr_object, _default_repr_attr)
 
-    def split(self, *filters) -> Union[PrettyDict[K, V], Tuple[PrettyDict[K, V], ...]]:
+    def split(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
         raise NotImplementedError
 
-    def filter(self, *filters) -> Union[PrettyDict[K, V], Tuple[PrettyDict[K, V], ...]]:
+    def filter(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
         raise NotImplementedError
 
-    def merge(self, *states) -> PrettyDict[K, V]:
+    def merge(self, *states) -> 'PrettyDict[K, V]':
         raise NotImplementedError
 
-    def subset(self, *filters) -> Union[PrettyDict[K, V], Tuple[PrettyDict[K, V], ...]]:
+    def subset(self, *filters) -> Union['PrettyDict[K, V]', Tuple['PrettyDict[K, V]', ...]]:
         """
-        Subset a ``PrettyDict`` into one or more ``PrettyDict``'s. The user must pass at least one
-        ``Filter`` (i.e. :class:`State`), and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the ``PrettyDict``).
+        Subset a :class:`PrettyDict` into one or more :class:`PrettyDict`'s. The user must pass at least one
+        `:class:`Filter` (i.e. :class:`State`), and the filters must be exhaustive (i.e. they must cover all
+        :class:`State` types in the :class:`PrettyDict`).
         """
         return self.filter(*filters)
 
 
+class NestedStateRepr(PrettyRepr):
+    def __init__(self, state: PrettyDict):
+        self.state = state
+
+    def __pretty_repr__(self):
+        yield PrettyType('', value_sep=': ', start='{', end='}')
+
+        for r in self.state.__pretty_repr__():
+            if isinstance(r, PrettyType):
+                continue
+            yield r
+
+    def __treescope_repr__(self, path, subtree_renderer):
+        children = {}
+        for k, v in self.state.items():
+            if isinstance(v, PrettyDict):
+                v = NestedStateRepr(v)
+            children[k] = v
+        # Render as the dictionary itself at the same path.
+        return subtree_renderer(children, path=path)
+
+
 def _default_repr_object(node: PrettyDict):
     yield PrettyType('', value_sep=': ', start='{', end='}')
 
@@ -375,20 +313,20 @@ def _default_repr_attr(node):
 
 class NestedDict(PrettyDict):
     """
-    A pytree-like structure that contains a ``Mapping`` from strings or integers to leaves.
+    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
 
     A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or nested
-    ``NestedDict`` and ``FlattedDict``.
+    :class:`NestedDict` and :class:`FlattedDict`.
     """
     __module__ = 'brainstate.util'
 
-    def __or__(self, other: NestedDict[K, V]) -> NestedDict[K, V]:
+    def __or__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
         if not other:
             return self
         assert isinstance(other, NestedDict), f'expected NestedDict; got {type(other).__qualname__}'
         return NestedDict.merge(self, other)
 
-    def __sub__(self, other: NestedDict[K, V]) -> NestedDict[K, V]:
+    def __sub__(self, other: 'NestedDict[K, V]') -> 'NestedDict[K, V]':
         if not other:
             return self
 
@@ -398,25 +336,25 @@ class NestedDict(PrettyDict):
         diff = {k: v for k, v in self_flat.items() if k not in other_flat}
         return NestedDict.from_flat(diff)
 
-    def to_flat(self) -> FlattedDict:
+    def to_flat(self) -> 'FlattedDict':
         """
         Flatten the nested mapping into a flat mapping.
 
         Returns:
-          The flattened mapping.
+            The flattened mapping.
         """
         return flat_mapping(self)
 
     @classmethod
-    def from_flat(cls, flat_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]]) -> NestedDict:
+    def from_flat(cls, flat_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]]) -> 'NestedDict':
         """
-        Create a ``NestedDict`` from a flat mapping.
+        Create a :class:`NestedDict` from a flat mapping.
 
         Args:
           flat_dict: The flat mapping.
 
         Returns:
-          The ``NestedDict``.
+          The :class:`NestedDict`.
         """
         nested_state = nest_mapping(dict(flat_dict))
         return cls(nested_state)
@@ -426,12 +364,12 @@ class NestedDict(PrettyDict):
         first: Filter,
         /,
         *filters: Filter
-    ) -> Union[NestedDict[K, V], Tuple[NestedDict[K, V], ...]]:
+    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
         """
-        Split a ``NestedDict`` into one or more ``NestedDict``'s. The
-        user must pass at least one ``Filter`` (i.e. :class:`State`),
+        Split a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
         and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the ``NestedDict``).
+        :class:`State` types in the :class:`NestedDict`).
 
         Example usage::
 
@@ -474,10 +412,10 @@ class NestedDict(PrettyDict):
         first: Filter,
         /,
         *filters: Filter,
-    ) -> Union[NestedDict[K, V], Tuple[NestedDict[K, V], ...]]:
+    ) -> Union['NestedDict[K, V]', Tuple['NestedDict[K, V]', ...]]:
         """
-        Filter a ``NestedDict`` into one or more ``NestedDict``'s. The
-        user must pass at least one ``Filter`` (i.e. :class:`State`).
+        Filter a :class:`NestedDict` into one or more :class:`NestedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
         This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
         except the filters can be non-exhaustive.
 
@@ -498,21 +436,21 @@ class NestedDict(PrettyDict):
 
     @staticmethod
     def merge(
-        state: NestedDict[K, V] | FlattedDict[K, V],
+        state: Union['NestedDict[K, V]', 'FlattedDict[K, V]'],
         /,
-        *states: NestedDict[K, V] | FlattedDict[K, V]
-    ) -> NestedDict[K, V]:
+        *states: Union['NestedDict[K, V]', 'FlattedDict[K, V]']
+    ) -> 'NestedDict[K, V]':
         """
         The inverse of :meth:`split()`.
 
-        ``merge`` takes one or more ``PrettyDict``'s and creates a new ``PrettyDict``.
+        ``merge`` takes one or more :class:`PrettyDict`'s and creates a new :class:`PrettyDict`.
 
         Args:
-          state: A ``PrettyDict`` object.
-          *states: Additional ``PrettyDict`` objects.
+          state: A :class:`PrettyDict` object.
+          *states: Additional :class:`PrettyDict` objects.
 
         Returns:
-          The merged ``PrettyDict``.
+          The merged :class:`PrettyDict`.
         """
         if not states:
             return state
@@ -548,11 +486,11 @@ class NestedDict(PrettyDict):
 
 class FlattedDict(PrettyDict):
     """
-    A pytree-like structure that contains a ``Mapping`` from strings or integers to leaves.
+    A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves.
 
     A valid leaf type is either :class:`State`, ``jax.Array``, ``numpy.ndarray`` or Python variables.
 
-    A ``NestedDict`` can be generated by either calling :func:`states()` or
+    A :class:`NestedDict` can be generated by either calling :func:`states()` or
     :func:`nodes()` on the :class:`Module`.
 
     Example usage::
@@ -631,13 +569,13 @@ class FlattedDict(PrettyDict):
     """
     __module__ = 'brainstate.util'
 
-    def __or__(self, other: FlattedDict[K, V]) -> FlattedDict[K, V]:
+    def __or__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
         if not other:
             return self
         assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
         return FlattedDict.merge(self, other)
 
-    def __sub__(self, other: FlattedDict[K, V]) -> FlattedDict[K, V]:
+    def __sub__(self, other: 'FlattedDict[K, V]') -> 'FlattedDict[K, V]':
         if not other:
             return self
         assert isinstance(other, FlattedDict), f'expected NestedDict; got {type(other).__qualname__}'
@@ -649,22 +587,22 @@ class FlattedDict(PrettyDict):
         Unflatten the flat mapping into a nested mapping.
 
         Returns:
-          The nested mapping.
+            The nested mapping.
         """
         return nest_mapping(self)
 
     @classmethod
     def from_nest(
         cls, nested_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]],
-    ) -> FlattedDict:
+    ) -> 'FlattedDict':
         """
-        Create a ``NestedDict`` from a flat mapping.
+        Create a :class:`NestedDict` from a flat mapping.
 
         Args:
           nested_dict: The flat mapping.
 
         Returns:
-          The ``NestedDict``.
+          The :class:`NestedDict`.
         """
         return flat_mapping(nested_dict)
 
@@ -673,19 +611,19 @@ class FlattedDict(PrettyDict):
         first: Filter,
         /,
         *filters: Filter
-    ) -> Union[FlattedDict[K, V], tuple[FlattedDict[K, V], ...]]:
+    ) -> Union['FlattedDict[K, V]', tuple['FlattedDict[K, V]', ...]]:
         """
-        Split a ``FlattedDict`` into one or more ``FlattedDict``'s. The
-        user must pass at least one ``Filter`` (i.e. :class:`State`),
+        Split a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`),
         and the filters must be exhaustive (i.e. they must cover all
-        :class:`State` types in the ``NestedDict``).
+        :class:`State` types in the :class:`NestedDict`).
 
         Arguments:
-          first: The first filter
-          *filters: The optional, additional filters to group the state into mutually exclusive substates.
+            first: The first filter
+            *filters: The optional, additional filters to group the state into mutually exclusive substates.
 
         Returns:
-          One or more ``States`` equal to the number of filters passed.
+            One or more ``States`` equal to the number of filters passed.
         """
         filters = (first, *filters)
         *states_, rest = _split_flatted_mapping(self, *filters)
@@ -705,19 +643,19 @@ class FlattedDict(PrettyDict):
         first: Filter,
         /,
         *filters: Filter,
-    ) -> Union[FlattedDict[K, V], Tuple[FlattedDict[K, V], ...]]:
+    ) -> Union['FlattedDict[K, V]', Tuple['FlattedDict[K, V]', ...]]:
         """
-        Filter a ``FlattedDict`` into one or more ``FlattedDict``'s. The
-        user must pass at least one ``Filter`` (i.e. :class:`State`).
+        Filter a :class:`FlattedDict` into one or more :class:`FlattedDict`'s. The
+        user must pass at least one `:class:`Filter` (i.e. :class:`State`).
         This method is similar to :meth:`split() <flax.nnx.NestedDict.state.split>`,
         except the filters can be non-exhaustive.
 
         Arguments:
-          first: The first filter
-          *filters: The optional, additional filters to group the state into mutually exclusive substates.
+            first: The first filter
+            *filters: The optional, additional filters to group the state into mutually exclusive substates.
 
         Returns:
-          One or more ``States`` equal to the number of filters passed.
+            One or more ``States`` equal to the number of filters passed.
         """
         *states_, _rest = _split_flatted_mapping(self, first, *filters)
         assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}'
@@ -729,21 +667,21 @@ class FlattedDict(PrettyDict):
 
     @staticmethod
     def merge(
-        state: FlattedDict[K, V] | NestedDict[K, V],
+        state: Union['FlattedDict[K, V]', 'NestedDict[K, V]'],
         /,
-        *states: FlattedDict[K, V] | NestedDict[K, V]
-    ) -> FlattedDict[K, V]:
+        *states: Union['FlattedDict[K, V]', 'NestedDict[K, V]']
+    ) -> 'FlattedDict[K, V]':
         """
         The inverse of :meth:`split()`.
 
-        ``merge`` takes one or more ``FlattedDict``'s and creates a new ``FlattedDict``.
+        ``merge`` takes one or more :class:`FlattedDict`'s and creates a new :class:`FlattedDict`.
 
         Args:
-          state: A ``PrettyDict`` object.
-          *states: Additional ``PrettyDict`` objects.
+          state: A :class:`PrettyDict` object.
+          *states: Additional :class:`PrettyDict` objects.
 
         Returns:
-          The merged ``PrettyDict``.
+          The merged :class:`PrettyDict`.
         """
         if not states:
             return state
@@ -759,14 +697,67 @@ class FlattedDict(PrettyDict):
         return FlattedDict(new_state)
 
     def to_dict_values(self):
+        """
+        Convert a FlattedDict containing State objects to a plain dictionary of values.
+        
+        This method extracts the underlying values from any State objects in the FlattedDict,
+        creating a new dictionary with the same keys but where each State object is replaced
+        by its value attribute. Non-State objects are kept as is.
+        
+        Returns:
+            dict: A dictionary with the same keys as the FlattedDict, but where each State
+                  object is replaced by its value attribute. Non-State objects remain unchanged.
+        
+        Example:
+            >>> flat_dict = FlattedDict({('model', 'layer1', 'weight'): ParamState(value=jnp.ones((10, 5)))})
+            >>> flat_dict.to_dict_values()
+            {('model', 'layer1', 'weight'): Array([[1., 1., ...]], dtype=float32)}
+        """
         from brainstate._state import State
-        return {k: v.value if isinstance(v, State) else v for k, v in self.items()}
+        return {
+            k: v.value if isinstance(v, State) else v
+            for k, v in self.items()
+        }
+
+    def assign_dict_values(self, data: dict):
+        """
+        Assign values from a dictionary to this FlattedDict.
+        
+        This method updates the values in the FlattedDict with values from the provided
+        dictionary. For keys that correspond to State objects, the value attribute of
+        the State is updated. For other keys, the value in the FlattedDict is directly
+        replaced with the new value.
+        
+        The method requires that all keys in the FlattedDict exist in the provided
+        dictionary, otherwise a KeyError is raised.
+        
+        Args:
+            data (dict): A dictionary containing the values to assign, where keys 
+                         must match those in the FlattedDict.
+        
+        Raises:
+            KeyError: If a key in the FlattedDict is not present in the provided dictionary.
+        
+        Example:
+            >>> flat_dict = FlattedDict({('model', 'weight'): ParamState(value=jnp.zeros((5, 5)))})
+            >>> flat_dict.assign_dict_values({('model', 'weight'): jnp.ones((5, 5))})
+            # The ParamState's value is now an array of ones
+        """
+        from brainstate._state import State
+        for k in self.keys():
+            if k not in data:
+                raise KeyError(f'Invalid key: {k!r}')
+            val = self[k]
+            if isinstance(val, State):
+                val.value = data[k]
+            else:
+                self[k] = data[k]
 
 
 def _split_nested_mapping(
-    mapping: NestedDict[K, V],
+    mapping: 'NestedDict[K, V]',
     *filters: Filter,
-) -> Tuple[NestedDict[K, V], ...]:
+) -> Tuple['NestedDict[K, V]', ...]:
     # check if the filters are exhaustive
     for i, filter_ in enumerate(filters):
         if filter_ in (..., True) and i != len(filters) - 1:
@@ -828,7 +819,7 @@ def _split_flatted_mapping(
     return tuple(FlattedDict(flat_state) for flat_state in flat_states)
 
 
-# register ``NestedDict`` as a pytree
+# register :class:`NestedDict` as a pytree
 def _nest_flatten_with_keys(x: NestedDict):
     items = sorted(x.items())
     children = tuple((jax.tree_util.DictKey(key), value) for key, value in items)
@@ -847,7 +838,7 @@ jax.tree_util.register_pytree_with_keys(NestedDict,
                                         _nest_unflatten)  # type: ignore[arg-type]
 
 
-# register ``FlattedDict`` as a pytree
+# register :class:`FlattedDict` as a pytree
 
 def _flat_unflatten(
     static: Tuple[K, ...],
@@ -892,3 +883,63 @@ def _list_repr_attr(node: PrettyList):
 
 def _list_repr_object(node: PrettyDict):
     yield PrettyType('', value_sep='', start='[', end=']')
+
+
+def _repr_object_general(node: PrettyDict):
+    """
+    Generate a general representation of a PrettyDict object.
+
+    This function is used to create a pretty representation of a PrettyDict
+    object, which includes the type of the object and its value separator.
+
+    Args:
+        node (PrettyDict): The PrettyDict object to be represented.
+
+    Yields:
+        PrettyType: A PrettyType object representing the type of the node,
+        with specified value separator, start, and end characters.
+    """
+    yield PrettyType(type(node), value_sep='=', start='(', end=')')
+
+
+def _repr_attribute_general(node):
+    """
+    Generate a pretty representation of the attributes of a node.
+
+    This function iterates over the attributes of a given node and attempts
+    to generate a pretty representation for each attribute. It handles
+    conversion of lists and dictionaries to their pretty representation
+    counterparts and yields a PrettyAttr object for each attribute.
+
+    Args:
+        node: The object whose attributes are to be represented.
+
+    Yields:
+        PrettyAttr: A PrettyAttr object representing the key and value of
+        each attribute in a pretty format.
+    """
+    for k, v in vars(node).items():
+        try:
+            res = node.__pretty_repr_item__(k, v)
+            if res is None:
+                continue
+            k, v = res
+        except AttributeError:
+            pass
+
+        if k is None:
+            continue
+
+        # convert list to PrettyList
+        if isinstance(v, list):
+            v = PrettyList(v)
+
+        # convert dict to PrettyDict
+        if isinstance(v, dict):
+            v = PrettyDict(v)
+
+        # convert PrettyDict to NestedStateRepr
+        if isinstance(v, PrettyDict):
+            v = NestedStateRepr(v)
+
+        yield PrettyAttr(k, v)

brainstate/util/_pretty_repr.py

@@ -15,8 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import dataclasses
 import threading
 from abc import ABC, abstractmethod

brainstate/util/_pretty_table.py

@@ -469,7 +469,7 @@ class PrettyTable:
         else:
             raise AttributeError(name)
 
-    def __getitem__(self, index: int | slice) -> PrettyTable:
+    def __getitem__(self, index: int | slice) -> 'PrettyTable':
         new = PrettyTable()
         new.field_names = self.field_names
         for attr in self._options:
@@ -2713,10 +2713,64 @@ class PrettyTable:
 ##############################
 
 
+# def _str_block_width(val: str) -> int:
+#     import wcwidth  # type: ignore[import-untyped]
+#
+#     return wcwidth.wcswidth(_re.sub("", val))
+
 def _str_block_width(val: str) -> int:
-    import wcwidth  # type: ignore[import-untyped]
+    """Calculate the visual width of a string, accounting for wide Unicode characters.
+
+    This is a custom implementation to replace the wcwidth dependency.
+
+    Args:
+        val: The string to measure
+
+    Returns:
+        The visual width of the string when displayed in a monospace context
+    """
+    # Remove ANSI escape sequences
+    val = _re.sub("", val)
+
+    # Fast path for ASCII-only strings
+    if all(ord(c) < 128 for c in val):
+        return len(val)
 
-    return wcwidth.wcswidth(_re.sub("", val))
+    width = 0
+    for char in val:
+        width += _char_width(char)
+
+    return width
+
+
+def _char_width(char: str) -> int:
+    """Calculate the display width of a single character.
+
+    Args:
+        char: A single Unicode character
+
+    Returns:
+        0 for control characters
+        2 for wide characters (CJK, emoji, etc.)
+        1 for all other characters
+    """
+    code = ord(char)
+
+    # Control characters and empty space
+    if code == 0 or code == 0x034F or (0x200B <= code <= 0x200F) or code == 0x2028 or code == 0x2029 or (
+        0x202A <= code <= 0x202E) or (0x2060 <= code <= 0x2063):
+        return 0
+
+    # Wide characters: CJK, emoji, etc.
+    if (0x1100 <= code <= 0x115F) or (0x2329 <= code <= 0x232A) or (0x2E80 <= code <= 0x303E) or (
+        0x3040 <= code <= 0x4DBF) or (0x4E00 <= code <= 0x9FFF) or (0xA000 <= code <= 0xA4CF) or (
+        0xAC00 <= code <= 0xD7A3) or (0xF900 <= code <= 0xFAFF) or (0xFE10 <= code <= 0xFE19) or (
+        0xFE30 <= code <= 0xFE6F) or (0xFF00 <= code <= 0xFF60) or (0xFFE0 <= code <= 0xFFE6) or (
+        0x1F300 <= code <= 0x1F64F) or (0x1F900 <= code <= 0x1F9FF):
+        return 2
+
+    # Default single-width character
+    return 1
 
 
 ##############################

brainstate/util/_scaling.py

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

brainstate/util/_struct.py

@@ -17,8 +17,6 @@
 
 """Utilities for defining custom classes that can be used with jax transformations."""
 
-from __future__ import annotations
-
 import collections
 import dataclasses
 from collections.abc import Hashable, Mapping