feat: redesign pretty repr for Node, Module, State and friends

brainstate/_state.py

@@ -50,7 +50,8 @@
 from brainstate._state_global_hooks import GlobalHookRegistry
 from brainstate._state_hook_manager import HookManager
 from brainstate.typing import ArrayLike, PyTree, Missing, Filter
-from brainstate.util import DictManager, PrettyObject, StateJaxTracer
+from brainstate.util import DictManager, PrettyObject, StateJaxTracer, PrettyType, PrettyAttr
+from brainstate.util._pretty_pytree import NestedStateRepr
 from brainstate.util.filter import Nothing
 
 __all__ = [
@@ -261,6 +262,17 @@ def __repr__(self) -> str:
 _INVALIDATED_TRACE_VALUE = _InvalidatedTraceValue()
 
 
+def _abstract_value(value: Any) -> Any:
+    """Abstract the array leaves of ``value`` for display in a ``repr``."""
+    try:
+        return jax.tree.map(shaped_abstractify, value)
+    except TypeError:
+        # e.g. the invalidated-trace sentinel, or any leaf that cannot be
+        # abstractified — show it raw rather than crash while building a repr
+        # (often inside an error message)
+        return value
+
+
 class State(Generic[A], PrettyObject):
     """
     A generic class representing a dynamic data pointer in the BrainState framework.
@@ -795,17 +807,54 @@ def to_state_ref(self: State[A]) -> TreefyState[A]:
         metadata.pop('_level', None)
         return TreefyState(type(self), self._value, **metadata)
 
+    def __pretty_repr__(self):
+        """Yield the pretty representation of the state.
+
+        The (abstracted) value is shown as the leading positional item,
+        followed by any non-empty metadata. A dict value without metadata is
+        shown in brace form.
+
+        Yields
+        ------
+        PrettyType or PrettyAttr
+            The type header followed by the items.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import jax.numpy as jnp
+            >>> import brainstate
+            >>> print(repr(brainstate.ParamState(jnp.zeros(3))))
+            ParamState(f32[3])
+            >>> print(repr(brainstate.ParamState({'w': jnp.zeros((2, 3)), 'b': jnp.zeros(3)})))
+            ParamState{b: f32[3], w: f32[2,3]}
+            >>> print(repr(brainstate.HiddenState(jnp.zeros(2), name='v')))
+            HiddenState(f32[2], name='v')
+        """
+        header, *items = super().__pretty_repr__()
+        if header.empty_repr == '...':
+            # already-seen object reported by ``yield_unique_pretty_repr_items``
+            yield header
+            return
+        value = [it for it in items if it.key == 'value']
+        others = [it for it in items if it.key != 'value']
+        if value and not others and isinstance(value[0].value, NestedStateRepr):
+            yield PrettyType(type(self), start='{', end='}', value_sep=': ')
+            for k, v in value[0].value.state.items():
+                key = k if isinstance(k, str) and k.isidentifier() else repr(k)
+                yield PrettyAttr(key, repr(v) if isinstance(v, str) else v)
+            return
+        yield header
+        for it in value:
+            yield PrettyAttr('', it.value)
+        yield from others
+
     def __pretty_repr_item__(self, k, v):
         if k in ['_level', '_source_info', '_been_writen', '_trace_state', '_hooks_manager']:
             return None
         if k == '_value':
-            try:
-                return 'value', jax.tree.map(shaped_abstractify, v)
-            except TypeError:
-                # e.g. the invalidated-trace sentinel, or any leaf that
-                # cannot be abstractified — show it raw rather than crash
-                # while building a repr (often inside an error message)
-                return 'value', v
+            return 'value', _abstract_value(v)
 
         if k == '_name':
             if self.name is None:
@@ -2656,13 +2705,12 @@ def __setattr__(self, name: str, value: Any) -> None: ...
         def __delattr__(self, name: str) -> None: ...
 
     def __pretty_repr_item__(self, k, v):
-        if k in ['_level', '_source_info', '_been_writen']:
+        if k in ['_level', '_source_info', '_been_writen', '_hooks_manager']:
             return None
-        if k == '_value':
-            return 'value', v
-
-        if k == '_name':
-            return None if v is None else ('name', v)
+        if k == 'value':
+            return 'value', _abstract_value(v)
+        if k in ('_name', 'tag'):
+            return None if v is None else (k.lstrip('_'), v)
         return k, v
 
     @property

brainstate/_state_test.py

@@ -2994,5 +2994,46 @@ def test_tree_set_value_correct_shape_updates(self):
         np.testing.assert_allclose(np.asarray(u.get_magnitude(got)), np.full((4, 3), 6.0))
 
 
+class TestStateReprFormat(unittest.TestCase):
+    """Compact ``State`` / ``TreefyState`` repr."""
+
+    def test_array_value_is_positional(self):
+        self.assertEqual(repr(brainstate.ParamState(jnp.zeros((2, 3)))), 'ParamState(f32[2,3])')
+        self.assertEqual(repr(brainstate.HiddenState(jnp.zeros(2, dtype=jnp.int32))), 'HiddenState(i32[2])')
+
+    def test_quantity_value(self):
+        self.assertEqual(repr(brainstate.HiddenState(jnp.zeros(3) * u.mV)), 'HiddenState(f32[3] mV)')
+
+    def test_dict_value_brace_form(self):
+        s = brainstate.ParamState({'weight': jnp.zeros((2, 3)), 'bias': jnp.zeros(3), 'a-b': jnp.zeros(1)})
+        self.assertEqual(repr(s), "ParamState{'a-b': f32[1], bias: f32[3], weight: f32[2,3]}")
+
+    def test_metadata_falls_back_to_parentheses(self):
+        s = brainstate.ParamState({'w': jnp.zeros(2)}, name='p', tag='t')
+        self.assertEqual(repr(s), "ParamState({'w': f32[2]}, name='p', tag={'t'})")
+
+    def test_scalar_and_python_values(self):
+        self.assertEqual(repr(brainstate.ShortTermState(jnp.asarray(1.5))), 'ShortTermState(f32[])')
+        self.assertEqual(repr(brainstate.ShortTermState(None)), 'ShortTermState(None)')
+
+    def test_treefy_state(self):
+        ref = brainstate.ParamState(jnp.zeros(3), name='w').to_state_ref()
+        text = repr(ref)
+        self.assertEqual(text, "TreefyState(type=ParamState, value=f32[3], name='w')")
+        self.assertNotIn('_hooks_manager', text)
+        self.assertNotIn('tag', text)
+
+    def test_repr_inside_jit(self):
+        texts = []
+
+        @jax.jit
+        def f(x):
+            texts.append(repr(brainstate.ShortTermState(x)))
+            return x
+
+        f(jnp.zeros(4))
+        self.assertEqual(texts, ['ShortTermState(f32[4])'])
+
+
 if __name__ == '__main__':
     unittest.main()

brainstate/graph/_node.py

@@ -17,16 +17,19 @@
 
 from __future__ import annotations
 
+import math
 from abc import ABCMeta
 from copy import deepcopy
 from typing import Any, Callable, TypeVar, TYPE_CHECKING
 
+import brainunit as u
 import jax
+import numpy as np
 
 from brainstate._error import TraceContextError
-from brainstate._state import State, TreefyState
+from brainstate._state import State, TreefyState, ParamState
 from brainstate.typing import Key
-from brainstate.util import PrettyObject
+from brainstate.util import PrettyObject, ReprSummary
 from ._operations import treefy_split, treefy_merge
 from ._walk import register_graph_node_type, iter_leaf
 
@@ -115,6 +118,43 @@ def __pretty_repr_item__(self, name: str, value: Any) -> tuple[str, Any] | None:
             return None
         return name, value
 
+    def __pretty_repr_summary__(self) -> ReprSummary:
+        """Summarise the size of the states reachable from this node.
+
+        Each :class:`~brainstate.State` is counted once even when it is shared
+        by several sub-nodes. Elements of :class:`~brainstate.ParamState`
+        values count as parameters; elements of every other state count as
+        state values. The raw stored value is inspected, so no state read
+        hooks are triggered.
+
+        Returns
+        -------
+        ReprSummary
+            Element counts and byte sizes, shown as ``# ...`` comments in the
+            node's ``repr``.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            >>> import brainstate
+            >>> s = brainstate.nn.Linear(3, 4).__pretty_repr_summary__()
+            >>> s.n_params, s.param_bytes
+            (16, 64)
+        """
+        n_params = param_bytes = n_values = value_bytes = 0
+        for _, leaf in iter_leaf(self):
+            if not isinstance(leaf, State):
+                continue
+            size, nbytes = _value_size(vars(leaf).get('_value'))
+            if isinstance(leaf, ParamState):
+                n_params += size
+                param_bytes += nbytes
+            else:
+                n_values += size
+                value_bytes += nbytes
+        return ReprSummary(n_params, param_bytes, n_values, value_bytes)
+
     def __deepcopy__(self: G, memo=None) -> G:
         graphdef, state = treefy_split(self)
         # ``state`` is a pytree of ``TreefyState`` whose only leaves are the
@@ -157,6 +197,19 @@ def check_valid_context(self, error_msg: Callable[[], str]) -> None:
 # -------------------------------
 
 
+def _value_size(value: Any) -> tuple[int, int]:
+    """Return ``(n_elements, n_bytes)`` over the array leaves of ``value``."""
+    size = nbytes = 0
+    for leaf in jax.tree.leaves(value):
+        leaf = u.get_mantissa(leaf)
+        if not hasattr(leaf, 'shape') or not hasattr(leaf, 'dtype'):
+            continue
+        n = math.prod(leaf.shape)
+        size += n
+        nbytes += n * np.dtype(leaf.dtype).itemsize
+    return size, nbytes
+
+
 def _node_flatten(node: Node) -> tuple[list[tuple[str, Any]], tuple[type]]:
     graph_invisible_attrs = getattr(node, 'graph_invisible_attrs', ())
     nodes = sorted(

brainstate/graph/_node_test.py

@@ -16,6 +16,7 @@
 import copy
 import unittest
 
+import brainunit as u
 import jax.numpy as jnp
 
 import brainstate
@@ -733,5 +734,75 @@ def __init__(self):
         self.assertIsNone(Outer().check_valid_context(lambda: 'ok'))
 
 
+class TestNodeReprSummary(unittest.TestCase):
+    """``__pretty_repr_summary__`` and the summary footer of ``Node`` reprs."""
+
+    def test_counts_params_and_other_states(self):
+        class Net(brainstate.graph.Node):
+            def __init__(self):
+                self.w = brainstate.ParamState(jnp.zeros((2, 3)))
+                self.b = brainstate.ParamState({'x': jnp.zeros(4, dtype=jnp.float16)})
+                self.h = brainstate.HiddenState(jnp.zeros(5, dtype=jnp.int32))
+
+        s = Net().__pretty_repr_summary__()
+        self.assertEqual((s.n_params, s.param_bytes), (10, 6 * 4 + 4 * 2))
+        self.assertEqual((s.n_state_values, s.state_bytes), (5, 20))
+
+    def test_shared_state_counted_once(self):
+        class Net(brainstate.graph.Node):
+            def __init__(self):
+                self.a = brainstate.ParamState(jnp.zeros(3))
+                self.b = self.a
+
+        self.assertEqual(Net().__pretty_repr_summary__().n_params, 3)
+
+    def test_quantity_and_non_array_leaves(self):
+        class Net(brainstate.graph.Node):
+            def __init__(self):
+                self.v = brainstate.HiddenState(jnp.zeros(4) * u.mV)
+                self.c = brainstate.ShortTermState(None)
+
+        s = Net().__pretty_repr_summary__()
+        self.assertEqual((s.n_state_values, s.state_bytes), (4, 16))
+
+    def test_summary_does_not_trigger_read_hooks(self):
+        class Net(brainstate.graph.Node):
+            def __init__(self):
+                self.w = brainstate.ParamState(jnp.zeros(3))
+
+        net = Net()
+        reads = []
+        net.w.register_hook('read', lambda ctx: reads.append(ctx))
+        repr(net)
+        self.assertEqual(reads, [])
+
+    def test_footer_in_repr(self):
+        class Net(brainstate.graph.Node):
+            def __init__(self):
+                self.w = brainstate.ParamState(jnp.zeros(3))
+
+        self.assertEqual(repr(Net()), 'Net(w=ParamState(f32[3]))  # 3 params (12 B)')
+        with brainstate.util.repr_options(summary=False):
+            self.assertEqual(repr(Net()), 'Net(w=ParamState(f32[3]))')
+
+    def test_stateless_node_has_no_footer(self):
+        class Empty(brainstate.graph.Node):
+            pass
+
+        self.assertEqual(repr(Empty()), 'Empty()')
+
+    def test_shared_and_cyclic_nodes(self):
+        class Leaf(brainstate.graph.Node):
+            pass
+
+        class Root(brainstate.graph.Node):
+            def __init__(self):
+                self.a = Leaf()
+                self.b = self.a
+                self.me = self
+
+        self.assertEqual(repr(Root()), 'Root(a=Leaf(), b=<shared: .a>, me=<cycle: .>)')
+
+
 if __name__ == '__main__':
     unittest.main()