From ab816a4267ed756a2c11b956afe7232bf987dafb Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Tue, 20 Oct 2026 16:59:23 +0800 Subject: [PATCH] feat: redesign pretty repr for Node, Module, State and friends Replace the one-attribute-per-line renderer with a width-aware compact tree layout shared by graph.Node, nn.Module, State, TreefyState and the util containers. - Layout: groups print on one line when they fit (default width 88), otherwise one item per line with trailing commas; [i] labels for sequences of objects; / markers. - Arrays summarised as f32[3,4], f32(0.5), f32[3] mV, key[]. - Summary: root footer "# 26 params (104 B)" and per-child comments, computed from unique states without triggering read hooks. - Options: ReprOptions, repr_options(), set_repr_options(), get_repr_options(); opt-in ANSI color for print() only (NO_COLOR, BRAINSTATE_REPR_COLOR); pretty_format() for arbitrary values. - Module: "(3,) -> (4,)" in/out header; State: positional value and ParamState{bias: ..., weight: ...} brace form. - Former hand-written reprs (elementwise activations, nn/interop transforms, HiData, RandomState, Mode family, ParamDescriber) now render through the engine. Fixes: TreefyState printed full arrays and _hooks_manager; Map/Param showed None attributes; Batching mislabelled batch_size as in_size; PermuteAxes printed as Transpose; HiData subclasses printed as HiData. Spec: docs/specs/pretty-repr-redesign.md --- brainstate/_state.py | 76 +- brainstate/_state_test.py | 41 + brainstate/graph/_node.py | 57 +- brainstate/graph/_node_test.py | 71 ++ brainstate/interop/_transforms.py | 38 +- brainstate/mixin.py | 58 +- brainstate/mixin_test.py | 14 + brainstate/nn/_common.py | 2 +- brainstate/nn/_common_test.py | 8 + brainstate/nn/_elementwise.py | 22 - brainstate/nn/_hidata.py | 128 +-- brainstate/nn/_hidata_test.py | 10 +- brainstate/nn/_module.py | 32 +- brainstate/nn/_module_test.py | 67 ++ brainstate/nn/_param.py | 5 +- brainstate/nn/_param_test.py | 5 + brainstate/nn/_transform.py | 95 ++- brainstate/random/_state.py | 26 +- brainstate/random/_state_test.py | 8 +- brainstate/util/__init__.py | 18 + brainstate/util/_pretty_pytree.py | 304 ++++++- brainstate/util/_pretty_repr.py | 862 ++++++++++++++++++-- brainstate/util/_pretty_repr_layout_test.py | 509 ++++++++++++ brainstate/util/_pretty_repr_test.py | 5 +- docs/specs/pretty-repr-redesign.md | 131 +++ 25 files changed, 2301 insertions(+), 291 deletions(-) create mode 100644 brainstate/util/_pretty_repr_layout_test.py create mode 100644 docs/specs/pretty-repr-redesign.md diff --git a/brainstate/_state.py b/brainstate/_state.py index 250975b9..a1bd196e 100644 --- a/brainstate/_state.py +++ b/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 diff --git a/brainstate/_state_test.py b/brainstate/_state_test.py index 10f39883..01a28fab 100644 --- a/brainstate/_state_test.py +++ b/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() diff --git a/brainstate/graph/_node.py b/brainstate/graph/_node.py index bd12c1ff..c32ede7a 100644 --- a/brainstate/graph/_node.py +++ b/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( diff --git a/brainstate/graph/_node_test.py b/brainstate/graph/_node_test.py index 06c946c8..326c2a05 100644 --- a/brainstate/graph/_node_test.py +++ b/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=, me=)') + + if __name__ == '__main__': unittest.main() diff --git a/brainstate/interop/_transforms.py b/brainstate/interop/_transforms.py index 2edf459e..69ef1dc9 100644 --- a/brainstate/interop/_transforms.py +++ b/brainstate/interop/_transforms.py @@ -31,6 +31,8 @@ import brainunit as u import numpy as np +from brainstate.util import PrettyAttr, PrettyRepr, PrettyType + __all__ = [ 'Transform', 'Identity', @@ -44,7 +46,7 @@ ] -class Transform: +class Transform(PrettyRepr): """Base class for an invertible array transform.""" def forward(self, x): @@ -55,8 +57,18 @@ def inverse(self, x): """Map a foreign-layout array back to the brainstate layout.""" raise NotImplementedError - def __repr__(self): - return f"{type(self).__name__}()" + def __pretty_repr__(self): + """Yield the pretty representation: the type and :meth:`_repr_items`.""" + yield PrettyType(type(self)) + for key, value in self._repr_items(): + yield PrettyAttr(key, repr(value) if isinstance(value, str) else value) + + def _repr_items(self): + """Return the ``(name, value)`` pairs shown in the ``repr``. + + An empty name renders the value positionally. + """ + return () class Identity(Transform): @@ -88,8 +100,8 @@ def forward(self, x): def inverse(self, x): return u.math.transpose(x, self.inv_perm) - def __repr__(self): - return f"Transpose(perm={self.perm})" + def _repr_items(self): + return [('perm', self.perm)] class PermuteAxes(Transpose): @@ -142,8 +154,8 @@ def forward(self, x): def inverse(self, x): return u.math.subtract(x, u.Quantity(self.constant, unit=u.get_unit(x))) - def __repr__(self): - return f"AddScalar({self.constant})" + def _repr_items(self): + return [('', self.constant)] class ReorderBlocks(Transform): @@ -177,8 +189,8 @@ def forward(self, x): def inverse(self, x): return self._apply(x, self.inv_order) - def __repr__(self): - return f"ReorderBlocks(axis={self.axis}, order={self.order})" + def _repr_items(self): + return [('axis', self.axis), ('order', self.order)] class SplitConcat(Transform): @@ -207,8 +219,8 @@ def forward(self, x) -> tuple: def inverse(self, xs) -> object: return u.math.concatenate(list(xs), axis=self.axis) - def __repr__(self): - return f"SplitConcat(axis={self.axis}, sizes={self.sizes})" + def _repr_items(self): + return [('axis', self.axis), ('sizes', self.sizes)] class Chain(Transform): @@ -233,5 +245,5 @@ def inverse(self, x): x = tf.inverse(x) return x - def __repr__(self): - return f"Chain({', '.join(map(repr, self.transforms))})" + def _repr_items(self): + return [('', t) for t in self.transforms] diff --git a/brainstate/mixin.py b/brainstate/mixin.py index 6ef228b7..b5a91fd9 100644 --- a/brainstate/mixin.py +++ b/brainstate/mixin.py @@ -28,6 +28,8 @@ import jax +from brainstate.util import PrettyAttr, PrettyRepr, PrettyType, pretty_format + __all__ = [ 'Mixin', 'ParamDesc', @@ -396,8 +398,8 @@ def __repr__(self): str A string showing the class and stored parameters. """ - args_str = ', '.join(repr(a) for a in self.args) - kwargs_str = ', '.join(f'{k}={v!r}' for k, v in self.kwargs.items()) + args_str = ', '.join(pretty_format(a) for a in self.args) + kwargs_str = ', '.join(f'{k}={pretty_format(v)}' for k, v in self.kwargs.items()) all_params = ', '.join(filter(None, [args_str, kwargs_str])) return f'ParamDescriber({self.cls.__name__}({all_params}))' @@ -917,7 +919,7 @@ def __getattr__(name): ) -class Mode(Mixin): +class Mode(Mixin, PrettyRepr): """ Base class for computation behavior modes. @@ -985,16 +987,16 @@ class Mode(Mixin): ... return x # No noise during evaluation """ - def __repr__(self): + def __pretty_repr__(self): """ - String representation of the mode. + Yield the pretty representation of the mode: its bare class name. - Returns - ------- - str - The class name of the mode. + Yields + ------ + PrettyType + The type header, without parentheses. """ - return self.__class__.__name__ + yield PrettyType(type(self), start='', end='') def _eq_key(self) -> Hashable: """ @@ -1148,7 +1150,7 @@ class JointMode(Mode): >>> >>> # Combine them >>> joint = brainstate.mixin.JointMode(training, batching) - >>> print(joint) # JointMode(Training, Batching(in_size=32, axis=0)) + >>> print(joint) # JointMode(Training, Batching(batch_size=32, batch_axis=0)) >>> >>> # Check if specific modes are present >>> assert joint.has(brainstate.mixin.Training) @@ -1201,16 +1203,18 @@ def __init__(self, *modes: Mode): def _eq_key(self) -> Hashable: return frozenset(self.modes) - def __repr__(self): + def __pretty_repr__(self): """ - String representation showing all combined modes. + Yield the pretty representation showing all combined modes. - Returns - ------- - str - A string showing the joint mode and its components. + Yields + ------ + PrettyType or PrettyAttr + The type header followed by one positional item per mode. """ - return f'{self.__class__.__name__}({", ".join([repr(m) for m in self.modes])})' + yield PrettyType(type(self)) + for mode in self.modes: + yield PrettyAttr('', mode) def has(self, mode: type): """ @@ -1346,7 +1350,7 @@ class Batching(Mode): >>> >>> # Create a batching mode >>> batching = brainstate.mixin.Batching(batch_size=32, batch_axis=0) - >>> print(batching) # Batching(in_size=32, axis=0) + >>> print(batching) # Batching(batch_size=32, batch_axis=0) >>> >>> # Access batch parameters >>> print(f"Processing {batching.batch_size} samples at once") @@ -1408,16 +1412,18 @@ def __init__(self, batch_size: int = 1, batch_axis: int = 0): def _eq_key(self) -> Hashable: return (self.batch_size, self.batch_axis) - def __repr__(self): + def __pretty_repr__(self): """ - String representation showing batch configuration. + Yield the pretty representation showing the batch configuration. - Returns - ------- - str - A string showing the batch size and axis. + Yields + ------ + PrettyType or PrettyAttr + The type header followed by ``batch_size`` and ``batch_axis``. """ - return f'{self.__class__.__name__}(in_size={self.batch_size}, axis={self.batch_axis})' + yield PrettyType(type(self)) + yield PrettyAttr('batch_size', self.batch_size) + yield PrettyAttr('batch_axis', self.batch_axis) class Training(Mode): diff --git a/brainstate/mixin_test.py b/brainstate/mixin_test.py index b73ddb5a..66e42ba1 100644 --- a/brainstate/mixin_test.py +++ b/brainstate/mixin_test.py @@ -499,6 +499,20 @@ def test_batching_repr(self): self.assertIn("64", repr(batching)) self.assertIn("0", repr(batching)) + def test_batching_repr_uses_constructor_names(self): + """The repr labels fields with the constructor argument names.""" + batching = brainstate.mixin.Batching(batch_size=64, batch_axis=1) + self.assertEqual(repr(batching), 'Batching(batch_size=64, batch_axis=1)') + + def test_mode_reprs_render_with_engine(self): + """Modes render through the pretty engine, nested modes included.""" + joint = brainstate.mixin.JointMode( + brainstate.mixin.Training(), brainstate.mixin.Batching(batch_size=8) + ) + self.assertEqual(repr(brainstate.mixin.Training()), 'Training') + self.assertEqual(repr(joint), 'JointMode(Training, Batching(batch_size=8, batch_axis=0))') + self.assertEqual(str(joint), repr(joint)) + def test_batching_is_mode(self): """Test Batching is a Mode.""" batching = brainstate.mixin.Batching() diff --git a/brainstate/nn/_common.py b/brainstate/nn/_common.py index 8e435c4b..07bc7050 100644 --- a/brainstate/nn/_common.py +++ b/brainstate/nn/_common.py @@ -503,7 +503,7 @@ def __pretty_repr_item__(self, name, value): '_call_state_axes' ]: return None - return name, value + return super().__pretty_repr_item__(name, value) def _integrate_state_axes(self, call_state_axes): if call_state_axes is None: diff --git a/brainstate/nn/_common_test.py b/brainstate/nn/_common_test.py index 6b9c6b81..b8f13665 100644 --- a/brainstate/nn/_common_test.py +++ b/brainstate/nn/_common_test.py @@ -400,6 +400,14 @@ def test_pretty_repr_hides_internal_fields(self): self.assertNotIn('dict_vmap_states', text) self.assertNotIn('_call_state_axes', text) + def test_pretty_repr_item_drops_none_and_strips_underscore(self): + """Map's item hook builds on Module's: None hidden, ``_x`` shown as ``x``.""" + m = Map(_ParamModule(), init_map_size=4) + self.assertIsNone(m.__pretty_repr_item__('spmd_axis_name', None)) + self.assertEqual(m.__pretty_repr_item__('_in_size', (3,)), ('in_size', (3,))) + self.assertIsNone(m.__pretty_repr_item__('_init', True)) + self.assertNotIn('=None', repr(m)) + def test_init_all_states_invalid_behavior_defensive(self): """A corrupted behavior triggers the defensive guard in init_all_states.""" m = Map(_ParamModule(), init_map_size=2) diff --git a/brainstate/nn/_elementwise.py b/brainstate/nn/_elementwise.py index fdd69076..80dc7f6c 100644 --- a/brainstate/nn/_elementwise.py +++ b/brainstate/nn/_elementwise.py @@ -96,8 +96,6 @@ def __call__(self, x: ArrayLike) -> ArrayLike: jnp.asarray(value, dtype=dtype)) return u.maybe_decimal(out * unit) - def __repr__(self): - return f'{self.__class__.__name__}(threshold={self.threshold}, value={self.value})' class ReLU(ElementWiseBlock): @@ -139,8 +137,6 @@ class ReLU(ElementWiseBlock): def __call__(self, x: ArrayLike) -> ArrayLike: return F.relu(x) - def __repr__(self): - return f'{self.__class__.__name__}()' class RReLU(ElementWiseBlock): @@ -664,8 +660,6 @@ def __init__(self, dim: int = -1) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.glu(x, self.dim) - def __repr__(self): - return f'{self.__class__.__name__}(dim={self.dim})' class GELU(ElementWiseBlock): @@ -710,8 +704,6 @@ def __init__(self, approximate: bool = False) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.gelu(x, approximate=self.approximate) - def __repr__(self): - return f'{self.__class__.__name__}(approximate={self.approximate})' class Hardshrink(ElementWiseBlock): @@ -757,8 +749,6 @@ def __init__(self, lambd: float = 0.5) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.hard_shrink(x, self.lambd) - def __repr__(self): - return f'{self.__class__.__name__}(lambd={self.lambd})' class LeakyReLU(ElementWiseBlock): @@ -808,8 +798,6 @@ def __init__(self, negative_slope: float = 1e-2) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.leaky_relu(x, self.negative_slope) - def __repr__(self): - return f'{self.__class__.__name__}(negative_slope={self.negative_slope})' class LogSigmoid(ElementWiseBlock): @@ -914,8 +902,6 @@ def __init__(self, lambd: float = 0.5) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.soft_shrink(x, self.lambd) - def __repr__(self): - return f'{self.__class__.__name__}(lambd={self.lambd})' class PReLU(ElementWiseBlock): @@ -987,8 +973,6 @@ def __init__(self, num_parameters: int = 1, init: float = 0.25, dtype=None) -> N def __call__(self, x: ArrayLike) -> ArrayLike: return F.prelu(x, self.weight.value) - def __repr__(self): - return f'{self.__class__.__name__}(num_parameters={self.num_parameters})' class Softsign(ElementWiseBlock): @@ -1093,8 +1077,6 @@ def __init__(self, dim: Optional[int] = None) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.softmin(x, -1 if self.dim is None else self.dim) - def __repr__(self): - return f'{self.__class__.__name__}(dim={self.dim})' class Softmax(ElementWiseBlock): @@ -1154,8 +1136,6 @@ def __init__(self, dim: Optional[int] = None) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.softmax(x, -1 if self.dim is None else self.dim) - def __repr__(self): - return f'{self.__class__.__name__}(dim={self.dim})' class Softmax2d(ElementWiseBlock): @@ -1237,8 +1217,6 @@ def __init__(self, dim: Optional[int] = None) -> None: def __call__(self, x: ArrayLike) -> ArrayLike: return F.log_softmax(x, -1 if self.dim is None else self.dim) - def __repr__(self): - return f'{self.__class__.__name__}(dim={self.dim})' class Identity(ElementWiseBlock): diff --git a/brainstate/nn/_hidata.py b/brainstate/nn/_hidata.py index 62ce2513..776777c0 100644 --- a/brainstate/nn/_hidata.py +++ b/brainstate/nn/_hidata.py @@ -10,6 +10,7 @@ from typing import Any, Dict import brainstate +from brainstate.util import PrettyAttr, PrettyRepr, PrettyType, format_array_like, pretty_repr_object from brainstate.util.struct import dataclass, field Array = brainstate.typing.ArrayLike @@ -24,7 +25,7 @@ def is_dataclass(cls): @dataclass -class HiData: +class HiData(PrettyRepr): """ Hierarchical state container for composed dynamics. @@ -49,11 +50,7 @@ class HiData: >>> data = HiData(name='config', learning_rate=0.01, batch_size=32) >>> print(data) - ParamData( - name='config', - learning_rate=0.01, - batch_size=32 - ) + HiData(name='config', learning_rate=0.01, batch_size=32) Create nested Data objects: @@ -62,17 +59,10 @@ class HiData: >>> model = HiData(name='model', weights=np.array([1, 2, 3])) >>> config = HiData(name='config', optimizer=optimizer, model=model) >>> print(config) - ParamData( + HiData( name='config', - optimizer=ParamData( - name='optimizer', - lr=0.001, - momentum=0.9 - ), - model=ParamData( - name='model', - weights=Array(shape=(3,), dtype=int64) - ) + optimizer=HiData(name='optimizer', lr=0.001, momentum=0.9), + model=HiData(name='model', weights=i64[3]), ) Access children using attribute or dictionary syntax: @@ -129,92 +119,36 @@ def values(self): return self.children.values() def __repr__(self) -> str: - """Return hierarchical string representation.""" - return self._repr_recursive(indent=0) - - def _repr_recursive(self, indent: int = 0) -> str: - """ - Generate hierarchical representation with indentation. - - Format: - HiData( - name='value', - child1=value1, - child2=HiData( - name='nested', - subchild=42 - ), - child3=value3 - ) - - Parameters - ---------- - indent - Current indentation level. - - Returns - ------- - String representation of this HiData and its children. - """ - indent_str = " " * indent - name_str = f"'{self.name}'" if self.name else "None" - - # All children use '=' separator - separator = "=" - - if not self.children: - # Empty Data object - return f"{indent_str}HiData(name={name_str})" - - # Start with Data( and name parameter - lines = [f"{indent_str}HiData("] - lines.append(f"{indent_str} name={name_str},") - - # Add children as parameters - child_items = list(self.children.items()) - for i, (key, value) in enumerate(child_items): - is_last = (i == len(child_items) - 1) - comma = "" if is_last else "," - - if isinstance(value, HiData): - # Recursively format nested Data objects - nested_repr = value._repr_recursive(indent + 1) - # Remove the leading indent from nested_repr since we're adding it ourselves - nested_lines = nested_repr.split('\n') - nested_lines[0] = nested_lines[0].lstrip() - nested_repr = '\n'.join(nested_lines) - lines.append(f"{indent_str} {key}{separator}{nested_repr}{comma}") - else: - # Format other values - value_repr = self._format_value(value) - lines.append(f"{indent_str} {key}{separator}{value_repr}{comma}") + """Return the hierarchical pretty representation.""" + return pretty_repr_object(self) - # Close with parenthesis - lines.append(f"{indent_str})") + __repr__._uses_pretty_repr = True - return "\n".join(lines) + def __pretty_repr__(self): + """Yield the pretty representation: ``name`` followed by the children. - def _format_value(self, value: Any) -> str: - """ - Format a non-HiData value for display. - - Parameters - ---------- - value - The value to format. + Nested ``HiData`` objects, containers and arrays are laid out by the + shared pretty-repr engine; other values whose ``repr`` is longer than + 60 characters are truncated with ``...``. - Returns - ------- - Formatted string representation. + Yields + ------ + PrettyType or PrettyAttr + The type header followed by the items. """ - if value is None: - return "None" - - # Handle arrays with shape information - if hasattr(value, 'shape') and hasattr(value, 'dtype'): - return f"Array(shape={value.shape}, dtype={value.dtype})" - - # Handle other types + yield PrettyType(type(self)) + yield PrettyAttr('name', repr(self.name)) + for key, value in self.children.items(): + yield PrettyAttr(key, self._format_value(value)) + + @staticmethod + def _format_value(value: Any) -> Any: + """Return ``value`` for the engine, or a truncated ``repr`` string.""" + if ( + isinstance(value, (PrettyRepr, list, tuple, dict)) + or format_array_like(value) is not None + ): + return value value_str = repr(value) if len(value_str) > 60: return f"{value_str[:57]}..." diff --git a/brainstate/nn/_hidata_test.py b/brainstate/nn/_hidata_test.py index 87119110..88453436 100644 --- a/brainstate/nn/_hidata_test.py +++ b/brainstate/nn/_hidata_test.py @@ -348,7 +348,7 @@ def test_array_formatting_numpy(self): # Root level uses '=' self.assertIn("name='arrays'", result) - self.assertIn("vector=Array(shape=(5,), dtype=int", result) + self.assertRegex(result, r"vector=i(32|64)\[5\]") def test_array_formatting_jax(self): """Test repr with JAX arrays.""" @@ -358,7 +358,7 @@ def test_array_formatting_jax(self): # Root level uses '=' self.assertIn("name='arrays'", result) - self.assertIn("vector=Array(shape=(3,), dtype=float", result) + self.assertIn("vector=f32[3]", result) def test_multidimensional_array_formatting(self): """Test repr with 2D arrays.""" @@ -367,7 +367,7 @@ def test_multidimensional_array_formatting(self): result = repr(data) # Root level uses '=' - self.assertIn("matrix=Array(shape=(3, 4), dtype=float", result) + self.assertRegex(result, r"matrix=f(32|64)\[3,4\]") def test_long_string_truncation(self): """Test that long strings are truncated.""" @@ -397,7 +397,7 @@ def test_mixed_children_repr(self): self.assertIn("name='mixed'", result) self.assertIn("number=42", result) self.assertIn("text='hello'", result) - self.assertIn("array=Array(shape=(3,), dtype=float", result) + self.assertRegex(result, r"array=f(32|64)\[3\]") self.assertIn("nested=", result) self.assertIn("none_val=None", result) self.assertIn("name='nested'", result) @@ -457,7 +457,7 @@ def test_repr_with_complex_nested_structure(self): self.assertIn("learning_rate=0.001", result) self.assertIn("layers=3", result) self.assertIn("epochs=100", result) - self.assertIn("weights=Array(shape=(10, 5), dtype=float", result) + self.assertRegex(result, r"weights=f(32|64)\[10,5\]") class TestModuleIsDataclassHelper(unittest.TestCase): diff --git a/brainstate/nn/_module.py b/brainstate/nn/_module.py index ecb42fee..6e5799c5 100644 --- a/brainstate/nn/_module.py +++ b/brainstate/nn/_module.py @@ -36,7 +36,7 @@ from brainstate.graph import Node, states, nodes, flatten from brainstate.mixin import ParamDescriber, ParamDesc, not_implemented, is_not_implemented from brainstate.typing import Size -from brainstate.util import FlattedDict, NestedDict +from brainstate.util import FlattedDict, NestedDict, PrettyAttr __all__ = [ 'Module', @@ -239,6 +239,36 @@ def __pretty_repr_item__(self, name, value): name = name[1:] # expose ``_name`` as ``name``, etc. return name, value + def __pretty_repr__(self): + """Yield the pretty representation of the module. + + Same as :meth:`brainstate.graph.Node.__pretty_repr__`, except that when + both ``in_size`` and ``out_size`` are shown they are merged into one + leading positional item ``in_size -> out_size``. + + Yields + ------ + PrettyType or PrettyAttr + The type header followed by the attribute items. + + Examples + -------- + .. code-block:: python + + >>> import brainstate + >>> print(repr(brainstate.nn.Dropout(0.5, in_size=(3,)))) + Dropout((3,) -> (3,), prob=0.5, broadcast_dims=()) + """ + items = list(super().__pretty_repr__()) + sizes = { + it.key: it for it in items[1:] + if isinstance(it, PrettyAttr) and it.key in ('in_size', 'out_size') + } + if len(sizes) == 2: + shape = PrettyAttr('', f"{sizes['in_size'].value!r} -> {sizes['out_size'].value!r}") + items = [items[0], shape] + [it for it in items[1:] if it.key not in sizes] + yield from items + def __call__(self, *args, **kwargs): return self.update(*args, **kwargs) diff --git a/brainstate/nn/_module_test.py b/brainstate/nn/_module_test.py index 7e83fd6c..02389ef6 100644 --- a/brainstate/nn/_module_test.py +++ b/brainstate/nn/_module_test.py @@ -1325,5 +1325,72 @@ def test_dropout_eval_is_identity(self): _testing.assert_allclose(out, x) +class TestModuleReprGolden(unittest.TestCase): + """Golden output of the compact ``Module`` repr.""" + + def test_sequential(self): + model = Sequential(brainstate.nn.Linear(3, 4), brainstate.nn.ReLU(), brainstate.nn.Linear(4, 2)) + self.assertEqual( + repr(model), + 'Sequential(\n' + ' (3,) -> (2,),\n' + ' layers=[\n' + ' [0] Linear((3,) -> (4,), weight=ParamState{bias: f32[4], weight: f32[3,4]}), # 16 params\n' + ' [1] ReLU(),\n' + ' [2] Linear((4,) -> (2,), weight=ParamState{bias: f32[2], weight: f32[4,2]}), # 10 params\n' + ' ],\n' + ') # 26 params (104 B)' + ) + + def test_shared_submodule_and_cycle(self): + class M(Module): + def __init__(self): + super().__init__() + self.a = brainstate.nn.Linear(2, 2) + self.b = self.a + self.me = self + + self.assertEqual( + repr(M()), + 'M(\n' + ' a=Linear((2,) -> (2,), weight=ParamState{bias: f32[2], weight: f32[2,2]}), # 6 params\n' + ' b=,\n' + ' me=,\n' + ') # 6 params (24 B)' + ) + + def test_in_out_size_merge_requires_both(self): + class OnlyIn(Module): + def __init__(self): + super().__init__() + self.in_size = (3,) + + self.assertEqual(repr(OnlyIn()), 'OnlyIn(in_size=(3,))') + + class Both(Module): + def __init__(self): + super().__init__() + self.in_size = (3,) + self.out_size = (5, 2) + + self.assertEqual(repr(Both()), 'Both((3,) -> (5, 2))') + + def test_named_module_quotes_name(self): + self.assertEqual(repr(Module(name='enc')), "Module(name='enc')") + + def test_max_depth_collapses_children(self): + model = Sequential(brainstate.nn.Linear(3, 4), brainstate.nn.ReLU()) + with brainstate.util.repr_options(max_depth=0): + self.assertEqual( + repr(model), + 'Sequential((3,) -> (4,), layers=[Linear(...), ReLU(...)]) # 16 params (64 B)' + ) + + def test_print_colored_but_repr_plain(self): + model = brainstate.nn.Linear(2, 2) + with brainstate.util.repr_options(color=True): + self.assertNotIn('\x1b[', repr(model)) + + if __name__ == '__main__': unittest.main() diff --git a/brainstate/nn/_param.py b/brainstate/nn/_param.py index 749c47a3..e6a03d48 100644 --- a/brainstate/nn/_param.py +++ b/brainstate/nn/_param.py @@ -397,9 +397,8 @@ def __pretty_repr_item__(self, name, value): 'precompute', '_cache_invalidation_hook_handle'): return None - if name.startswith('_'): - return None if value is None else (name[1:], value) # skip the first `_` - return name, value + # drops ``None`` / invisible attributes and exposes ``_x`` as ``x`` + return super().__pretty_repr_item__(name, value) @classmethod def init( diff --git a/brainstate/nn/_param_test.py b/brainstate/nn/_param_test.py index baa9fb54..b3fe1b84 100644 --- a/brainstate/nn/_param_test.py +++ b/brainstate/nn/_param_test.py @@ -1904,6 +1904,11 @@ def test_public_field_passthrough(self): param = Param(jnp.array([1.0])) self.assertEqual(param.__pretty_repr_item__('fit', True), ('fit', True)) + def test_public_none_field_dropped(self): + """A public field with a None value is dropped, as for any Node.""" + param = Param(jnp.array([1.0])) + self.assertIsNone(param.__pretty_repr_item__('reg', None)) + def test_private_none_field_dropped(self): """A leading-underscore field with a None value is dropped.""" param = Param(jnp.array([1.0])) diff --git a/brainstate/nn/_transform.py b/brainstate/nn/_transform.py index 94fa28bc..ed721bd2 100644 --- a/brainstate/nn/_transform.py +++ b/brainstate/nn/_transform.py @@ -15,7 +15,7 @@ from abc import ABC, abstractmethod -from typing import Sequence +from typing import Any, Sequence, Tuple import brainunit as u import jax @@ -23,6 +23,8 @@ from brainstate.typing import ArrayLike from jax import Array +from brainstate.util import PrettyAttr, PrettyRepr, PrettyType + __all__ = [ 'Transform', 'IdentityT', @@ -86,7 +88,7 @@ def save_exp(x, max_value: float = 20.0): return u.math.exp(x) -class Transform(ABC): +class Transform(PrettyRepr, ABC): r""" Abstract base class for bijective parameter transformations. @@ -126,9 +128,24 @@ class Transform(ABC): """ __module__ = 'brainstate.nn' - def __repr__(self) -> str: - """Return a string representation of the transform.""" - return f"{self.__class__.__name__}()" + def __pretty_repr__(self): + """Yield the pretty representation: the type and :meth:`_repr_items`.""" + yield PrettyType(type(self)) + for key, value in self._repr_items(): + yield PrettyAttr(key, repr(value) if isinstance(value, str) else value) + + def _repr_items(self) -> Sequence[Tuple[str, Any]]: + """Return the ``(name, value)`` pairs shown in the ``repr``. + + Subclasses override this to report their constructor arguments; an + empty name renders the value positionally. + + Returns + ------- + sequence of tuple of (str, Any) + The items to display, in order. + """ + return () def __call__(self, x: ArrayLike) -> Array: r""" @@ -250,8 +267,6 @@ class IdentityT(Transform): """Identity transformation (no-op).""" __module__ = 'brainstate.nn' - def __repr__(self): - return "IdentityT()" def forward(self, x: ArrayLike) -> Array: return x @@ -344,8 +359,8 @@ def __init__(self, lower: ArrayLike, upper: ArrayLike) -> None: self.width = upper - lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"SigmoidT(lower={self.lower}, upper={self.lower + self.width})" + def _repr_items(self): + return [('lower', self.lower), ('upper', self.lower + self.width)] def forward(self, x: ArrayLike) -> Array: r""" @@ -481,8 +496,8 @@ def __init__(self, lower: ArrayLike) -> None: self.lower = lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"SoftplusT(lower={self.lower})" + def _repr_items(self): + return [('lower', self.lower)] def forward(self, x: ArrayLike) -> Array: """ @@ -615,8 +630,8 @@ def __init__(self, upper: ArrayLike) -> None: """ super().__init__(upper) - def __repr__(self) -> str: - return f"NegSoftplusT(upper={self.lower})" + def _repr_items(self): + return [('upper', self.lower)] def forward(self, x: ArrayLike) -> Array: """ @@ -698,8 +713,8 @@ def __init__(self, lower: ArrayLike) -> None: self.lower = lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"LogT(lower={self.lower})" + def _repr_items(self): + return [('lower', self.lower)] def forward(self, x: ArrayLike) -> Array: return self.lower + save_exp(x) * self.unit @@ -735,8 +750,8 @@ def __init__(self, lower: ArrayLike) -> None: self.lower = lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"ExpT(lower={self.lower})" + def _repr_items(self): + return [('lower', self.lower)] def forward(self, x: ArrayLike) -> Array: return self.lower + save_exp(x) * self.unit @@ -764,8 +779,8 @@ def __init__(self, lower: ArrayLike, upper: ArrayLike) -> None: self.width = upper - lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"TanhT(lower={self.lower}, upper={self.lower + self.width})" + def _repr_items(self): + return [('lower', self.lower), ('upper', self.lower + self.width)] def forward(self, x: ArrayLike) -> Array: return self.lower + self.width * (jnp.tanh(x) + 1.0) / 2.0 @@ -790,8 +805,8 @@ def __init__(self, lower: ArrayLike, upper: ArrayLike) -> None: self.width = upper - lower self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"SoftsignT(lower={self.lower}, upper={self.lower + self.width})" + def _repr_items(self): + return [('lower', self.lower), ('upper', self.lower + self.width)] def forward(self, x: ArrayLike) -> Array: return self.lower + self.width * (x / (1.0 + u.math.abs(x)) + 1.0) / 2.0 @@ -863,8 +878,8 @@ def __init__(self, lower: ArrayLike, upper: ArrayLike) -> None: self.lower = lower self.upper = upper - def __repr__(self) -> str: - return f"ClipT(lower={self.lower}, upper={self.upper})" + def _repr_items(self): + return [('lower', self.lower), ('upper', self.upper)] def forward(self, x: ArrayLike) -> Array: """ @@ -997,8 +1012,8 @@ def __init__(self, scale: ArrayLike, shift: ArrayLike): self.a = scale self.b = shift - def __repr__(self) -> str: - return f"AffineT(scale={self.a}, shift={self.b})" + def _repr_items(self): + return [('scale', self.a), ('shift', self.b)] def forward(self, x: ArrayLike) -> Array: """ @@ -1128,9 +1143,8 @@ def __init__(self, *transforms: Transform) -> None: super().__init__() self.transforms: Sequence[Transform] = transforms - def __repr__(self) -> str: - transforms_str = ", ".join(repr(t) for t in self.transforms) - return f"ChainT({transforms_str})" + def _repr_items(self): + return [('', t) for t in self.transforms] def forward(self, x: ArrayLike) -> Array: """ @@ -1286,8 +1300,8 @@ def __init__(self, mask: ArrayLike, transform: Transform, safe_value: ArrayLike self.transform = transform self.safe_value = safe_value - def __repr__(self) -> str: - return f"MaskedT(mask=..., transform={repr(self.transform)})" + def _repr_items(self): + return [('mask', self.mask), ('transform', self.transform)] def forward(self, x: ArrayLike) -> Array: """ @@ -1356,6 +1370,9 @@ class ReluT(Transform): def __init__(self, lower_bound: float = 0.0): self.lower_bound = lower_bound + def _repr_items(self): + return [('lower_bound', self.lower_bound)] + def forward(self, x: ArrayLike) -> Array: return u.math.relu(x) + self.lower_bound @@ -1404,8 +1421,6 @@ def __init__(self) -> None: """Initialize the positive transformation.""" super().__init__() - def __repr__(self) -> str: - return "PositiveT()" def forward(self, x: ArrayLike) -> Array: """Transform unbounded input to positive values.""" @@ -1450,8 +1465,6 @@ def __init__(self) -> None: """Initialize the negative transformation.""" super().__init__() - def __repr__(self) -> str: - return "NegativeT()" def forward(self, x: ArrayLike) -> Array: """Transform unbounded input to negative values.""" @@ -1529,8 +1542,8 @@ def __init__(self, lower: ArrayLike, upper: ArrayLike, beta: float = 1.0) -> Non self.beta = beta self.unit = u.get_unit(lower) - def __repr__(self) -> str: - return f"ScaledSigmoidT(lower={self.lower}, upper={self.lower + self.width}, beta={self.beta})" + def _repr_items(self): + return [('lower', self.lower), ('upper', self.lower + self.width), ('beta', self.beta)] def forward(self, x: ArrayLike) -> Array: """Transform unbounded input to bounded interval.""" @@ -1601,8 +1614,8 @@ def __init__(self, lmbda: float = 0.5) -> None: super().__init__() self.lmbda = lmbda - def __repr__(self) -> str: - return f"PowerT(lmbda={self.lmbda})" + def _repr_items(self): + return [('lmbda', self.lmbda)] def forward(self, x: ArrayLike) -> Array: """Apply the power transformation.""" @@ -1647,8 +1660,6 @@ def __init__(self) -> None: """Initialize the ordered transformation.""" super().__init__() - def __repr__(self) -> str: - return "OrderedT()" def forward(self, x: ArrayLike) -> Array: """Transform unconstrained input to ordered vectors.""" @@ -1701,8 +1712,6 @@ def __init__(self) -> None: """Initialize the simplex transformation.""" super().__init__() - def __repr__(self) -> str: - return "SimplexT()" def forward(self, x: ArrayLike) -> Array: """Transform unconstrained input to simplex.""" @@ -1778,8 +1787,6 @@ def __init__(self) -> None: """Initialize the unit vector transformation.""" super().__init__() - def __repr__(self) -> str: - return "UnitVectorT()" def forward(self, x: ArrayLike) -> Array: """Project input onto unit sphere.""" diff --git a/brainstate/random/_state.py b/brainstate/random/_state.py index 02c1131a..e32e4875 100644 --- a/brainstate/random/_state.py +++ b/brainstate/random/_state.py @@ -29,6 +29,7 @@ from brainstate import environ from brainstate._compatible_import import Tracer, concrete_or_error from brainstate._state import State +from brainstate.util import PrettyAttr, PrettyType from brainstate.transform._named_scope import named_scope from brainstate.typing import DTypeLike, Size, SeedOrKey from ._impl import ( @@ -138,8 +139,29 @@ def __init__( self._backup = None - def __repr__(self): - return f'{self.__class__.__name__}({self.value})' + def __pretty_repr__(self): + """Yield the pretty representation: the key type and its raw bits. + + The stored key is inspected directly, so no read hooks run and a lazy + raw-key placeholder is not materialised. The key bits are shown when + they are concrete and short, since they identify the random stream. + + Examples + -------- + .. code-block:: python + + >>> import brainstate + >>> print(repr(brainstate.random.RandomState(42))) + RandomState(key[], data=[0, 42]) + """ + key = vars(self).get('_value') + yield PrettyType(type(self)) + yield PrettyAttr('', key) + data = jr.key_data(key) if _is_typed_key(key) else key + if isinstance(data, (np.ndarray, jax.Array)) and not isinstance(data, Tracer) and data.size <= 8: + yield PrettyAttr('data', np.asarray(data).tolist()) + if self.name is not None: + yield PrettyAttr('name', repr(self.name)) def check_if_deleted(self): # Lazily wrap a stored raw-key placeholder into a typed key on first read. diff --git a/brainstate/random/_state_test.py b/brainstate/random/_state_test.py index d03a63e9..633ec56d 100644 --- a/brainstate/random/_state_test.py +++ b/brainstate/random/_state_test.py @@ -757,10 +757,12 @@ def tearDown(self): TRACE_CONTEXT.state_stack.pop() def test_repr_contains_value(self): - """The repr embeds the class name and the underlying key value.""" + """The repr embeds the class name, the key type and the raw key bits.""" text = repr(self.rs) - self.assertIn('RandomState', text) - self.assertIn(str(self.rs.value), text) + self.assertIn('RandomState(key<', text) + bits = np.asarray(jax.random.key_data(self.rs.value)).tolist() + self.assertIn(f'data={bits}', text) + self.assertEqual(str(self.rs), text) def test_clone_is_independent(self): """Cloning yields a distinct instance with a distinct key.""" diff --git a/brainstate/util/__init__.py b/brainstate/util/__init__.py index 5efc208b..f62b4ecd 100644 --- a/brainstate/util/__init__.py +++ b/brainstate/util/__init__.py @@ -49,6 +49,15 @@ PrettyRepr, PrettyMapping, MappingReprMixin, + pretty_repr_object, + pretty_format, + ReprOptions, + ReprSummary, + get_repr_options, + set_repr_options, + repr_options, + format_array_like, + format_bytes, ) # Import tracer utilities from ._tracers import ( @@ -98,6 +107,15 @@ def breakpoint_if(*args, **kwargs): 'PrettyAttr', 'PrettyRepr', 'PrettyMapping', + 'pretty_repr_object', + 'pretty_format', + 'ReprOptions', + 'ReprSummary', + 'get_repr_options', + 'set_repr_options', + 'repr_options', + 'format_array_like', + 'format_bytes', 'MappingReprMixin', # Other utilities diff --git a/brainstate/util/_pretty_pytree.py b/brainstate/util/_pretty_pytree.py index 39e4888a..64fc262a 100644 --- a/brainstate/util/_pretty_pytree.py +++ b/brainstate/util/_pretty_pytree.py @@ -41,7 +41,7 @@ >>> state = NestedDict({'layer1': {'weight': 1, 'bias': 2}}) >>> flat = state.to_flat() >>> print(flat) - FlattedDict({('layer1', 'weight'): 1, ('layer1', 'bias'): 2}) + {('layer1', 'weight'): 1, ('layer1', 'bias'): 2} """ from collections import abc @@ -91,7 +91,8 @@ class PrettyObject(PrettyRepr): ... def __init__(self, value): ... self.value = value >>> tree = MyTree(42) - >>> print(tree) # Uses __pretty_repr__ for display + >>> print(tree) + MyTree(value=42) """ def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]: @@ -138,6 +139,12 @@ def __pretty_repr_item__(self, k: Any, v: Any) -> Optional[Tuple[Any, Any]]: PrettyReprTree = PrettyObject +def _quote_str(v: Any) -> Any: + # ``PrettyAttr`` inserts ``str`` values verbatim; quote string *data* so it + # is distinguishable from identifiers (``name='w'`` rather than ``name=w``). + return repr(v) if isinstance(v, str) else v + + # The empty node is a struct.dataclass to be compatible with JAX. @dataclass class _EmptyNode: @@ -189,15 +196,15 @@ def flat_mapping( >>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}} >>> flat_xs = flat_mapping(xs) >>> flat_xs - FlattedDict({('foo',): 1, ('bar', 'a'): 2}) + {('foo',): 1, ('bar', 'a'): 2} >>> # With separator >>> flat_mapping(xs, sep='/') - FlattedDict({'foo': 1, 'bar/a': 2}) + {'foo': 1, 'bar/a': 2} >>> # Keep empty nodes >>> flat_mapping(xs, keep_empty_nodes=True) - FlattedDict({('foo',): 1, ('bar', 'a'): 2, ('bar', 'b'): _EmptyNode()}) + {('foo',): 1, ('bar', 'a'): 2, ('bar', 'b'): _EmptyNode()} Notes ----- @@ -267,12 +274,12 @@ def nest_mapping( ... } >>> xs = nest_mapping(flat_xs) >>> xs - NestedDict({'foo': 1, 'bar': {'a': 2}}) + {'foo': 1, 'bar': {'a': 2}} >>> # With separator >>> flat_xs_str = {'foo': 1, 'bar/a': 2} >>> nest_mapping(flat_xs_str, sep='/') - NestedDict({'foo': 1, 'bar': {'a': 2}}) + {'foo': 1, 'bar': {'a': 2}} See Also -------- @@ -422,6 +429,8 @@ def __repr__(self) -> str: """ return pretty_repr_object(self) + __repr__._uses_pretty_repr = True + def __pretty_repr__(self) -> Generator[Union[PrettyType, PrettyAttr], None, None]: """Generate pretty representation items for this dictionary. @@ -606,7 +615,7 @@ def _default_repr_attr(node: PrettyDict) -> Generator[PrettyAttr, None, None]: if isinstance(v, PrettyDict): v = NestedStateRepr(v) - yield PrettyAttr(repr(k), v) + yield PrettyAttr(repr(k), _quote_str(v)) class NestedDict(PrettyDict): @@ -634,7 +643,11 @@ class NestedDict(PrettyDict): ... }) >>> flat = state.to_flat() >>> print(flat) - FlattedDict({('layer1', 'weight'): ..., ('layer1', 'bias'): ..., ...}) + { + ('layer1', 'weight'): f32[3,3], + ('layer1', 'bias'): f32[3], + ('layer2', 'weight'): f32[3,1], + } See Also -------- @@ -854,6 +867,273 @@ def replace_by_pure_dict( self.update(nest_mapping(current_flat)) +class FlattedDict(PrettyDict): + """ + 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 :class:`NestedDict` can be generated by either calling :func:`states()` or + :func:`nodes()` on the :class:`Module`. + + Example usage:: + + >>> import brainstate as brainstate + >>> import jax.numpy as jnp + >>> + >>> class Model(brainstate.nn.Module): + ... def __init__(self): + ... super().__init__() + ... self.batchnorm = brainstate.nn.BatchNorm1d([10, 3]) + ... self.linear = brainstate.nn.Linear(2, 3) + ... def __call__(self, x): + ... return self.linear(self.batchnorm(x)) + >>> + >>> model = Model() + + >>> # retrieve the states of the model + >>> model.states() # with the same to the function of ``brainstate.graph.states()`` + { + ('batchnorm', 'running_mean'): BatchState(f32[1,3]), + ('batchnorm', 'running_var'): BatchState(f32[1,3]), + ('batchnorm', 'weight'): NormalizationParamState{bias: f32[1,3], scale: f32[1,3]}, + ('linear', 'weight'): ParamState{bias: f32[3], weight: f32[2,3]}, + } + + >>> # retrieve the nodes of the model + >>> model.nodes() # with the same to the function of ``brainstate.graph.nodes()`` + { + ('batchnorm',): BatchNorm1d( + (10, 3) -> (10, 3), + affine=True, + bias_initializer=Constant(value=0.0), + scale_initializer=Constant(value=1.0), + dtype=float32, + track_running_stats=True, + momentum=f32(0.99), + epsilon=f32(1e-05), + use_fast_variance=True, + feature_axes=(1,), + running_mean=BatchState(f32[1,3]), + running_var=BatchState(f32[1,3]), + weight=NormalizationParamState{bias: f32[1,3], scale: f32[1,3]}, + ), # 6 params + ('linear',): Linear((2,) -> (3,), weight=ParamState{bias: f32[3], weight: f32[2,3]}), # 9 params + (): Model(batchnorm=, linear=), # 15 params + } + """ + __module__ = 'brainstate.util' + + def __or__(self, other: 'NestedDict') -> 'NestedDict': + if not other: + return self + if not isinstance(other, NestedDict): + raise TypeError(f'expected NestedDict; got {type(other).__qualname__}') + return NestedDict.merge(self, other) + + def __sub__(self, other: 'NestedDict') -> 'NestedDict': + if not other: + return self + + if not isinstance(other, NestedDict): + raise TypeError(f'expected NestedDict; got {type(other).__qualname__}') + self_flat = self.to_flat() + other_flat = other.to_flat() + 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': + """ + Flatten the nested mapping into a flat mapping. + + Returns + ------- + The flattened mapping. + """ + return flat_mapping(self) + + @classmethod + def from_flat(cls, flat_dict: abc.Mapping[PathParts, V] | Iterable[tuple[PathParts, V]]) -> 'NestedDict': + """ + Create a :class:`NestedDict` from a flat mapping. + + Parameters + ---------- + flat_dict + The flat mapping. + + Returns + ------- + The :class:`NestedDict`. + """ + nested_state = nest_mapping(dict(flat_dict)) + return cls(nested_state) + + def split(self, *filters: Filter) -> Union['NestedDict', Tuple['NestedDict', ...]]: + """ + 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 :class:`NestedDict`). + + Example usage:: + + >>> import brainstate as brainstate + + >>> class Model(brainstate.nn.Module): + ... def __init__(self): + ... super().__init__() + ... self.batchnorm = brainstate.nn.BatchNorm1d([10, 3]) + ... self.linear = brainstate.nn.Linear(2, 3) + ... def __call__(self, x): + ... return self.linear(self.batchnorm(x)) + + >>> model = Model() + >>> state_map = brainstate.graph.treefy_states(model) + >>> param, others = state_map.treefy_split(brainstate.ParamState, ...) + + Parameters + ---------- + 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. + """ + *states_, rest = _split_nested_mapping(self, *filters) + if rest: + raise ValueError(f'Non-exhaustive filters, got a non-empty remainder: {rest}.\n' + f'Use `...` to match all remaining elements.') + + states: NestedDict | Tuple[NestedDict, ...] + if len(states_) == 1: + states = states_[0] + else: + states = tuple(states_) + return states # type: ignore[bad-return-type] + + def filter(self, *filters: Filter) -> Union['NestedDict', Tuple['NestedDict', ...]]: + """ + 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() `, + except the filters can be non-exhaustive. + + Parameters + ---------- + 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. + """ + states_ = _split_nested_mapping(self, *filters) + assert len(states_) == len(filters) + 1, f'Expected {len(filters) + 1} states, got {len(states_)}' + if len(filters) == 1: + return states_[0] + else: + return tuple(states_[:-1]) + + @staticmethod + def merge(*states) -> 'NestedDict': + """ + The inverse of :meth:`split()`. + + ``merge`` takes one or more :class:`PrettyDict`'s and creates a new :class:`PrettyDict`. + + Parameters + ---------- + *states + Additional :class:`PrettyDict` objects. + + Returns + ------- + The merged :class:`PrettyDict`. + """ + new_state: FlattedDict = FlattedDict() + for state in states: + if isinstance(state, NestedDict): + new_state.update(state.to_flat()) # type: ignore[attribute-error] # pytype is wrong here + elif isinstance(state, FlattedDict): + new_state.update(state) + else: + raise TypeError(f'Expected Nested or Flatted Mapping, got {type(state)} instead.') + return NestedDict.from_flat(new_state) + + def to_pure_dict(self) -> Dict[str, Any]: + """Convert to a pure nested dictionary structure. + + This method creates a standard Python dictionary with the same nested structure + as this NestedDict, without any special class wrappers. + + Returns + ------- + Dict[str, Any] + A pure nested dictionary representation. + + Examples + -------- + >>> nested = NestedDict({'a': {'b': 1, 'c': 2}}) + >>> pure = nested.to_pure_dict() + >>> type(pure) + + """ + # ``keep_empty_nodes=True`` so that empty nested mappings (e.g. ``{'a': {}}``) + # survive the flatten/nest round-trip; the default would silently drop them, + # contradicting this method's "same nested structure" contract. ``nest_mapping`` + # restores the empty-node sentinel back to an empty ``dict``. + flat_values = flat_mapping(self, keep_empty_nodes=True) + return nest_mapping(flat_values).to_dict() + + def replace_by_pure_dict( + self, + pure_dict: Dict[str, Any], + replace_fn: Optional[SetValueFn] = None + ) -> None: + """Replace values in this NestedDict using a pure dictionary. + + This method updates the values in this NestedDict with values from a standard + Python dictionary. For :class:`State` objects with a ``replace`` method, the + replace method is called; otherwise, values are directly assigned. + + Parameters + ---------- + pure_dict + A pure dictionary with matching structure containing new values. + replace_fn + Optional custom function to replace values. Takes ``(old_value, new_value)`` + and returns the updated value. Defaults to calling ``replace()`` method if + available, otherwise direct assignment. + + Raises + ------ + ValueError + If a key in ``pure_dict`` is not found in this NestedDict. + + Examples + -------- + >>> from brainstate._state import State + >>> nested = NestedDict({'a': State(1), 'b': 2}) + >>> nested.replace_by_pure_dict({'a': 10, 'b': 20}) + >>> nested['a'].value + 10 + """ + if replace_fn is None: + replace_fn = lambda x, v: x.replace(v) if hasattr(x, 'replace') else v + current_flat = self.to_flat() + for kp, v in flat_mapping(pure_dict).items(): + if kp not in current_flat: + raise ValueError(f'key in pure_dict not available in state: {kp}') + current_flat[kp] = replace_fn(current_flat[kp], v) + self.update(nest_mapping(current_flat)) + + class FlattedDict(PrettyDict): """ A pytree-like structure that contains a :class:`Mapping` from strings or integers to leaves. @@ -1385,6 +1665,8 @@ def __repr__(self) -> str: """ return pretty_repr_object(self) + __repr__._uses_pretty_repr = True + def tree_flatten(self) -> Tuple[list, Tuple]: """Flatten the list for JAX pytree operations. @@ -1438,7 +1720,7 @@ def _list_repr_attr(node: PrettyList) -> Generator[PrettyAttr, None, None]: v = PrettyDict(v) if isinstance(v, PrettyDict): v = NestedStateRepr(v) - yield PrettyAttr('', v) + yield PrettyAttr('', _quote_str(v)) def _list_repr_object(node: PrettyList) -> Generator[PrettyType, None, None]: @@ -1523,4 +1805,4 @@ def _repr_attribute_general(node: Any) -> Generator[PrettyAttr, None, None]: if isinstance(v, PrettyDict): v = NestedStateRepr(v) - yield PrettyAttr(k, v) + yield PrettyAttr(k, _quote_str(v)) diff --git a/brainstate/util/_pretty_repr.py b/brainstate/util/_pretty_repr.py index 891cea91..83004df0 100644 --- a/brainstate/util/_pretty_repr.py +++ b/brainstate/util/_pretty_repr.py @@ -22,12 +22,20 @@ objects, with support for nested structures and circular reference detection. """ +import contextlib import dataclasses +import math +import os import threading from abc import ABC, abstractmethod -from functools import partial from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional +import brainunit as u +import jax +import numpy as np + +from brainstate._compatible_import import Tracer + __all__ = [ 'yield_unique_pretty_repr_items', 'PrettyType', @@ -35,6 +43,15 @@ 'PrettyRepr', 'PrettyMapping', 'MappingReprMixin', + 'pretty_repr_object', + 'pretty_format', + 'ReprOptions', + 'ReprSummary', + 'get_repr_options', + 'set_repr_options', + 'repr_options', + 'format_array_like', + 'format_bytes', ] A = TypeVar('A') @@ -60,6 +77,9 @@ class PrettyType: The indentation for nested elements. empty_repr : str, default='' The representation for empty objects. + ref : str or None, default=None + When set, the object is rendered as this literal marker (for example + ``''``) instead of being expanded. """ type: Union[str, type] start: str = '(' @@ -67,6 +87,7 @@ class PrettyType: value_sep: str = '=' elem_indent: str = ' ' empty_repr: str = '' + ref: Optional[str] = None @dataclasses.dataclass @@ -77,18 +98,24 @@ class PrettyAttr: Attributes ---------- key : str - The attribute name or key. + The attribute name or key. An empty key renders the value as a + positional item (without ``key``). value : Union[str, Any] - The attribute value. + The attribute value. A ``str`` value is inserted verbatim; any other + value is rendered recursively. start : str, default='' Optional prefix for the attribute. end : str, default='' Optional suffix for the attribute. + comment : str or None, default=None + Optional trailing ``# comment``, shown only when the enclosing group is + broken over several lines. """ key: str value: Union[str, Any] start: str = '' end: str = '' + comment: Optional[str] = None class PrettyRepr(ABC): @@ -115,10 +142,7 @@ class PrettyRepr(ABC): ... >>> obj = MyObject('foo', 42) >>> print(obj) - MyObject{ - key=foo, - value=42 - } + MyObject{key=foo, value=42} """ __slots__ = () @@ -145,6 +169,26 @@ def __repr__(self) -> str: """ return pretty_repr_object(self) + def __str__(self) -> str: + """ + Generate the display string used by ``print``. + + Identical to ``repr`` except that ANSI colors are applied when enabled + through :func:`repr_options` / :func:`set_repr_options` (``color=True``) + and the ``NO_COLOR`` environment variable is not set. A subclass that + overrides ``__repr__`` with its own format gets that format unchanged. + + Returns + ------- + str + The formatted, possibly colored, string representation. + """ + if not getattr(type(self).__repr__, '_uses_pretty_repr', False): + return repr(self) + return pretty_repr_object(self, color=_color_enabled()) + + __repr__._uses_pretty_repr = True + def pretty_repr_elem(obj: PrettyType, elem: Any) -> str: """ @@ -183,56 +227,785 @@ def pretty_repr_elem(obj: PrettyType, elem: Any) -> str: return f'{obj.elem_indent}{elem.start}{elem.key}{obj.value_sep}{value}{elem.end}' -def pretty_repr_object(obj: PrettyRepr) -> str: +def pretty_repr_object(obj: PrettyRepr, *, color: bool = False) -> str: """ - Generates a pretty string representation of an object that implements the PrettyRepr interface. + Render an object implementing the :class:`PrettyRepr` interface. - This function utilizes the __pretty_repr__ method of the PrettyRepr interface to obtain - a structured representation of the object, which includes both the type and attributes - of the object in a human-readable format. + The items yielded by ``obj.__pretty_repr__()`` are turned into a document + tree which is then laid out to fit the configured line width (see + :class:`ReprOptions`): a group is printed on a single line when it fits, + and is otherwise broken one item per line with trailing commas. Repeated + and cyclic references are rendered as ```` / + ```` markers. Parameters ---------- obj : PrettyRepr - The object for which the pretty representation is to be generated. The object must - implement the PrettyRepr interface. + The object to render. + color : bool, default=False + Whether to apply ANSI colors to the output. Returns ------- str - A string that represents the object in a pretty format, including its type and attributes. - The format is determined by the PrettyType and PrettyAttr instances yielded by the - __pretty_repr__ method of the object. + The rendered representation. Raises ------ TypeError - If the provided object does not implement the PrettyRepr interface or if the first item - yielded by the __pretty_repr__ method is not an instance of PrettyType. + If ``obj`` does not implement :class:`PrettyRepr`, or if the first item + yielded by its ``__pretty_repr__`` method is not a :class:`PrettyType`. + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import PrettyMapping, pretty_repr_object, repr_options + >>> m = PrettyMapping({'a': 1, 'b': [1, 2]}, type_name='Cfg') + >>> print(pretty_repr_object(m)) + Cfg{'a': 1, 'b': [1, 2]} + >>> + >>> with repr_options(width=16): + ... print(pretty_repr_object(m)) + Cfg{ + 'a': 1, + 'b': [1, 2], + } """ if not isinstance(obj, PrettyRepr): raise TypeError(f'Object {obj!r} is not representable') - iterator = obj.__pretty_repr__() - obj_repr = next(iterator) + return _render(obj, color) + + +def pretty_format(value: Any, *, color: bool = False) -> str: + """ + Format any value with the pretty-repr engine. + + Unlike :func:`pretty_repr_object`, ``value`` does not need to implement + :class:`PrettyRepr`: lists, tuples, dicts, dataclasses and array-likes are + laid out and summarised the same way as inside a pretty ``repr``. + + Parameters + ---------- + value : Any + The value to format. + color : bool, default=False + Whether to apply ANSI colors to the output. + + Returns + ------- + str + The formatted value. + + Examples + -------- + .. code-block:: python + + >>> import numpy as np + >>> from brainstate.util import pretty_format + >>> pretty_format({'w': np.zeros((2, 3), dtype=np.float32), 'n': 3}) + "{'w': f32[2,3], 'n': 3}" + """ + return _render(value, color) + + +def _render(value: Any, color: bool) -> str: + options = get_repr_options() + is_root = CONTEXT.render_depth == 0 + CONTEXT.render_depth += 1 + try: + doc = _DocBuilder(options).build(value, '', 0) + text = _Layout(options, color).render(doc, 0, options.width) + if is_root and options.summary and isinstance(doc, _Group): + summary = _get_summary(value) + footer = '' if summary is None else summary.footer() + if footer: + text += _style(f' # {footer}', 'comment', color) + return text + finally: + CONTEXT.render_depth -= 1 + + +# --------------------------------------------------------------------------- +# Options +# --------------------------------------------------------------------------- - # repr object - if not isinstance(obj_repr, PrettyType): - raise TypeError(f'First item must be PrettyType, got {type(obj_repr).__name__}') - - # repr attributes - elem_reprs = tuple(map(partial(pretty_repr_elem, obj_repr), iterator)) - elems = ',\n'.join(elem_reprs) - if elems: - elems = '\n' + elems + '\n' - else: - elems = obj_repr.empty_repr - # repr object type - type_repr = obj_repr.type if isinstance(obj_repr.type, str) else obj_repr.type.__name__ +@dataclasses.dataclass(frozen=True) +class ReprOptions: + """ + Options controlling how pretty representations are rendered. + + Parameters + ---------- + width : int, default=88 + Target line width. Groups that fit within the remaining width are + printed on one line; others are broken one item per line. + max_depth : int or None, default=None + Number of nested object levels expanded below the root; deeper objects + are shown as ``Type(...)``. ``0`` collapses all children of the root. + ``None`` means unlimited. + max_items : int or None, default=64 + Maximum number of items shown per group; the rest are summarised as + ``... (+N more)``. ``None`` means unlimited. + summary : bool, default=True + Whether to show parameter / state summaries for graph nodes. + color : bool, default=False + Whether ``str()`` / ``print()`` apply ANSI colors. ``repr()`` is never + colored. Ignored when the ``NO_COLOR`` environment variable is set. + short_dtypes : bool, default=True + Whether arrays are shown with short dtype names (``f32[3]``) instead of + the full names (``float32[3]``). + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import ReprOptions + >>> ReprOptions(width=60).width + 60 + """ + width: int = 88 + max_depth: Optional[int] = None + max_items: Optional[int] = 64 + summary: bool = True + color: bool = False + short_dtypes: bool = True + + def __post_init__(self): + if not isinstance(self.width, int) or self.width < 1: + raise ValueError(f'width must be a positive int, got {self.width!r}') + if self.max_depth is not None and (not isinstance(self.max_depth, int) or self.max_depth < 0): + raise ValueError(f'max_depth must be None or a non-negative int, got {self.max_depth!r}') + if self.max_items is not None and (not isinstance(self.max_items, int) or self.max_items < 1): + raise ValueError(f'max_items must be None or a positive int, got {self.max_items!r}') + - # return repr - return f'{type_repr}{obj_repr.start}{elems}{obj_repr.end}' +def _env_flag(name: str) -> bool: + return os.environ.get(name, '').strip().lower() in ('1', 'true', 'yes', 'on') + + +_GLOBAL_OPTIONS = ReprOptions(color=_env_flag('BRAINSTATE_REPR_COLOR')) + + +def get_repr_options() -> ReprOptions: + """ + Return the pretty-repr options active in the current thread. + + Returns + ------- + ReprOptions + The innermost :func:`repr_options` override, or the process-wide + default set by :func:`set_repr_options`. + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import get_repr_options, repr_options + >>> with repr_options(width=40): + ... print(get_repr_options().width) + 40 + """ + options = CONTEXT.options + return _GLOBAL_OPTIONS if options is None else options + + +def set_repr_options(**kwargs) -> ReprOptions: + """ + Update the process-wide default pretty-repr options. + + Parameters + ---------- + **kwargs + Fields of :class:`ReprOptions` to change. + + Returns + ------- + ReprOptions + The previous default options, which can be restored with + ``set_repr_options(**dataclasses.asdict(previous))``. + + Raises + ------ + TypeError + If an unknown option name is given. + ValueError + If an option value is invalid. + + Examples + -------- + .. code-block:: python + + >>> import dataclasses + >>> from brainstate.util import set_repr_options, get_repr_options + >>> previous = set_repr_options(width=100) + >>> get_repr_options().width + 100 + >>> _ = set_repr_options(**dataclasses.asdict(previous)) + """ + global _GLOBAL_OPTIONS + previous = _GLOBAL_OPTIONS + _GLOBAL_OPTIONS = dataclasses.replace(previous, **kwargs) + return previous + + +@contextlib.contextmanager +def repr_options(**kwargs) -> Iterator[ReprOptions]: + """ + Temporarily override pretty-repr options in the current thread. + + Overrides nest, and the previous options are restored on exit, including + when an exception is raised. + + Parameters + ---------- + **kwargs + Fields of :class:`ReprOptions` to change. + + Yields + ------ + ReprOptions + The options in effect inside the ``with`` block. + + Raises + ------ + TypeError + If an unknown option name is given. + ValueError + If an option value is invalid. + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import PrettyMapping, repr_options + >>> m = PrettyMapping({'a': 1, 'b': 2}) + >>> with repr_options(width=8): + ... print(repr(m)) + { + 'a': 1, + 'b': 2, + } + """ + previous = CONTEXT.options + new = dataclasses.replace(get_repr_options(), **kwargs) + CONTEXT.options = new + try: + yield new + finally: + CONTEXT.options = previous + + +def _color_enabled() -> bool: + return get_repr_options().color and 'NO_COLOR' not in os.environ + + +# --------------------------------------------------------------------------- +# Summaries +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class ReprSummary: + """ + Size summary of a graph node, shown as a ``# ...`` comment in its repr. + + Parameters + ---------- + n_params : int + Number of parameter elements (leaves of ``ParamState`` values). + param_bytes : int + Memory footprint of the parameters in bytes. + n_state_values : int + Number of elements held by non-parameter states. + state_bytes : int + Memory footprint of the non-parameter states in bytes. + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import ReprSummary + >>> s = ReprSummary(n_params=26, param_bytes=104) + >>> s.footer() + '26 params (104 B)' + >>> s.comment() + '26 params' + """ + n_params: int = 0 + param_bytes: int = 0 + n_state_values: int = 0 + state_bytes: int = 0 + + def footer(self) -> str: + """Return the long form shown after the root object (empty if all zero).""" + parts = [] + if self.n_params: + parts.append(f'{_count(self.n_params, "param")} ({format_bytes(self.param_bytes)})') + if self.n_state_values: + parts.append(f'{_count(self.n_state_values, "state value")} ({format_bytes(self.state_bytes)})') + return ', '.join(parts) + + def comment(self) -> str: + """Return the short form shown after nested objects (empty if all zero).""" + if self.n_params: + return _count(self.n_params, 'param') + if self.n_state_values: + return _count(self.n_state_values, 'state value') + return '' + + +def _count(n: int, noun: str) -> str: + return f'{n:,} {noun}' if n == 1 else f'{n:,} {noun}s' + + +def format_bytes(n: int) -> str: + """ + Format a byte count with binary units. + + Parameters + ---------- + n : int + Number of bytes. + + Returns + ------- + str + Human-readable size such as ``'104 B'`` or ``'1.5 KiB'``. + + Examples + -------- + .. code-block:: python + + >>> from brainstate.util import format_bytes + >>> format_bytes(104) + '104 B' + >>> format_bytes(1536) + '1.5 KiB' + >>> format_bytes(3 * 1024 ** 3) + '3.0 GiB' + """ + if n < 1024: + return f'{int(n)} B' + size = float(n) + for unit in ('KiB', 'MiB', 'GiB'): + size /= 1024 + if size < 1024 or unit == 'GiB': + break + return f'{size:.1f} {unit}' + + +def _get_summary(obj: Any) -> Optional[ReprSummary]: + fn = getattr(obj, '__pretty_repr_summary__', None) + if fn is None: + return None + try: + return fn() + except Exception: + # A summary is decoration; never let it break a repr. + return None + + +# --------------------------------------------------------------------------- +# Array formatting +# --------------------------------------------------------------------------- + +_DTYPE_PREFIX = ( + ('bfloat', 'bf'), + ('float', 'f'), + ('uint', 'u'), + ('int', 'i'), + ('complex', 'c'), +) + + +def _short_dtype(name: str) -> str: + for prefix, short in _DTYPE_PREFIX: + if name.startswith(prefix) and name[len(prefix):].isdigit(): + return short + name[len(prefix):] + return name + + +def _is_array_like(x: Any) -> bool: + if isinstance(x, (np.ndarray, np.generic, jax.Array, jax.ShapeDtypeStruct)): + return True + # Abstract values such as ``ShapedArray``; duck-typed to stay robust + # across JAX versions. + return type(x).__name__ == 'ShapedArray' and hasattr(x, 'shape') and hasattr(x, 'dtype') + + +def _is_concrete(x: Any) -> bool: + if isinstance(x, (np.ndarray, np.generic)): + return True + return isinstance(x, jax.Array) and not isinstance(x, Tracer) + + +def format_array_like(x: Any, short_dtypes: Optional[bool] = None) -> Optional[str]: + """ + Summarise an array-like value as ``dtype[shape]``. + + Parameters + ---------- + x : Any + A ``jax.Array`` (including tracers), ``numpy.ndarray``, NumPy scalar, + ``jax.ShapeDtypeStruct``, ``ShapedArray``, or a ``brainunit.Quantity`` + wrapping one of them. + short_dtypes : bool or None, default=None + Whether to use short dtype names. ``None`` uses the active + :class:`ReprOptions`. + + Returns + ------- + str or None + The summary (``'f32[3,4]'``, ``'f32(0.5)'``, ``'f32[3] mV'``), or + ``None`` when ``x`` is not array-like. + + Examples + -------- + .. code-block:: python + + >>> import numpy as np + >>> import brainunit as u + >>> from brainstate.util import format_array_like + >>> format_array_like(np.zeros((3, 4), dtype=np.float32)) + 'f32[3,4]' + >>> format_array_like(np.zeros((3, 4), dtype=np.float32), short_dtypes=False) + 'float32[3,4]' + >>> format_array_like(np.float32(0.5)) + 'f32(0.5)' + >>> format_array_like(np.ones(3, dtype=np.float32) * u.mV) + 'f32[3] mV' + >>> format_array_like([1, 2]) is None + True + """ + if short_dtypes is None: + short_dtypes = get_repr_options().short_dtypes + if isinstance(x, u.Quantity): + mantissa = x.mantissa + if isinstance(mantissa, (int, float, complex)) and not isinstance(mantissa, bool): + return str(x) + if not _is_array_like(mantissa): + return None + if x.is_unitless: + return format_array_like(mantissa, short_dtypes) + if np.ndim(mantissa) == 0 and _is_concrete(mantissa): + return str(x) + return f'{format_array_like(mantissa, short_dtypes)} {x.unit}' + if not _is_array_like(x): + return None + try: + dtype = np.dtype(x.dtype).name + numpy_dtype = True + except TypeError: + # extended dtypes such as typed PRNG keys (``key``) + dtype = str(x.dtype) + numpy_dtype = False + if short_dtypes: + dtype = _short_dtype(dtype) + shape = tuple(x.shape) + if not shape and numpy_dtype and _is_concrete(x): + # NumPy scalars print the shortest round-tripping form for their own + # precision (``0.99`` rather than ``0.9900000095367432`` for float32). + return f'{dtype}({str(np.asarray(x)[()])})' + return f'{dtype}[{",".join(str(d) for d in shape)}]' + + +# --------------------------------------------------------------------------- +# Document tree and layout +# --------------------------------------------------------------------------- + +_ANSI = { + 'type': '\x1b[1;36m', + 'key': '\x1b[34m', + 'array': '\x1b[32m', + 'ref': '\x1b[2;35m', + 'comment': '\x1b[2m', +} +_ANSI_RESET = '\x1b[0m' + + +def _style(text: str, kind: Optional[str], color: bool) -> str: + if not color or kind is None or not text: + return text + return f'{_ANSI[kind]}{text}{_ANSI_RESET}' + + +@dataclasses.dataclass(eq=False) +class _Leaf: + text: str + kind: Optional[str] = None + + +@dataclasses.dataclass(eq=False) +class _Item: + key: str # attribute name / dict key; '' for positional items + sep: str # separator between key and value + value: Any # _Leaf | _Group + start: str = '' + end: str = '' + label: str = '' # e.g. '[0] ', shown only when the parent is broken + comment: Union[None, str, Callable[[], str]] = None + + def prefix_width(self) -> int: + return len(self.start) + (len(self.key) + len(self.sep) if self.key else 0) + + +@dataclasses.dataclass(eq=False) +class _Group: + head: str + open: str + close: str + items: list + empty: str = '' + indent: int = 2 + single_comma: bool = False # one-element tuple: ``(x,)`` when flat + + +def _is_marker(item: _Item) -> bool: + return isinstance(item.value, _Leaf) and item.value.kind == 'comment' + + +def _has_generated_repr(value: Any) -> bool: + """Whether ``value`` is a dataclass instance using the generated ``__repr__``.""" + if isinstance(value, type) or not dataclasses.is_dataclass(value): + return False + # ``dataclasses`` builds ``__repr__`` from source compiled as ```` + # and wraps it with ``reprlib.recursive_repr`` (which sets ``__wrapped__``); + # a hand-written ``__repr__`` has neither property. + wrapped = getattr(type(value).__repr__, '__wrapped__', None) + code = getattr(wrapped, '__code__', None) + return code is not None and code.co_filename == '' + + +def _close(iterator: Iterator) -> None: + """Close a partially consumed ``__pretty_repr__`` generator (runs its cleanup).""" + close = getattr(iterator, 'close', None) + if close is not None: + close() + + +def _child_path(path: str, key: str, sep: str, index: int) -> str: + if not key: + return f'{path}[{index}]' + if sep.strip() == '=' and key.isidentifier(): + return f'{path}.{key}' + return f'{path}[{key}]' + + +class _DocBuilder: + """Turns objects into a document tree, tracking shared/cyclic references.""" + + def __init__(self, options: ReprOptions): + self.options = options + # PrettyRepr objects already rendered: id -> (object, path). The object + # is kept alive so that ids of temporaries cannot be reused. + self.seen: dict[int, tuple[Any, str]] = {} + # objects currently being rendered (ancestors): id -> path + self.active: dict[int, str] = {} + + def build(self, value: Any, path: str, depth: int) -> Union[_Leaf, _Group]: + if isinstance(value, PrettyRepr): + return self._build_pretty(value, path, depth) + if isinstance(value, type): + return _Leaf(value.__name__) + if type(value) in (list, tuple, dict) or _has_generated_repr(value): + key = id(value) + if key in self.active: + return _Leaf(f'', 'ref') + self.active[key] = path + try: + return self._build_container(value, path, depth) + finally: + del self.active[key] + text = format_array_like(value, self.options.short_dtypes) + if text is not None: + return _Leaf(text, 'array') + return _Leaf(repr(value)) + + def _build_container(self, value: Any, path: str, depth: int) -> _Group: + if type(value) is list: + return self._build_items(_Group('', '[', ']', []), [('', '', v) for v in value], path, depth) + if type(value) is tuple: + group = _Group('', '(', ')', [], single_comma=len(value) == 1) + return self._build_items(group, [('', '', v) for v in value], path, depth) + if type(value) is dict: + entries = [(repr(k), ': ', v) for k, v in value.items()] + return self._build_items(_Group('', '{', '}', []), entries, path, depth) + entries = [ # dataclass with the generated ``__repr__`` + PrettyAttr(f.name, getattr(value, f.name)) + for f in dataclasses.fields(value) if f.repr + ] + group = _Group(type(value).__name__, '(', ')', []) + return self._build_items(group, entries, path, depth, '=', quote=True) + + def _build_items( + self, group: _Group, entries, path: str, depth: int, attr_sep: str = '', quote: bool = False + ) -> _Group: + """Build ``entries`` (``PrettyAttr``\\ s or ``(key, sep, value)`` triples) into ``group``. + + ``PrettyAttr`` string values are inserted verbatim unless ``quote``. + """ + limit = self.options.max_items + n_extra = 0 + for index, entry in enumerate(entries): + if limit is not None and index >= limit: + n_extra += 1 + continue + if isinstance(entry, PrettyAttr): + sep = attr_sep if entry.key else '' + item = _Item(entry.key, sep, None, entry.start, entry.end, comment=entry.comment) + raw = entry.value + if isinstance(raw, str) and not quote: + # protocol: string values are inserted verbatim + item.value = _Leaf(raw) + group.items.append(item) + continue + else: + key, sep, raw = entry + item = _Item(key, sep, None) + item.value = self.build(raw, _child_path(path, item.key, item.sep, index), depth) + if ( + item.comment is None + and self.options.summary + and isinstance(item.value, _Group) + and hasattr(raw, '__pretty_repr_summary__') + ): + item.comment = self._summary_comment(raw) + group.items.append(item) + if n_extra: + group.items.append(_Item('', '', _Leaf(f'... (+{n_extra} more)', 'comment'))) + self._label(group) + return group + + @staticmethod + def _summary_comment(raw: Any) -> Callable[[], str]: + def comment() -> str: + summary = _get_summary(raw) + return '' if summary is None else summary.comment() + + return comment + + @staticmethod + def _label(group: _Group): + """Add ``[i]`` labels to broken sequences of nested objects.""" + if group.open not in ('[', '(') or any(it.key for it in group.items): + return + if not any( + isinstance(it.value, _Group) and it.value.head + or isinstance(it.value, _Leaf) and it.value.kind == 'ref' + for it in group.items + ): + return + for i, it in enumerate(group.items): + if not _is_marker(it): + it.label = f'[{i}] ' + + def _build_pretty(self, obj: PrettyRepr, path: str, depth: int) -> Union[_Leaf, _Group]: + key = id(obj) + if key in self.active: + return _Leaf(f'', 'ref') + if key in self.seen: + return _Leaf(f'', 'ref') + + iterator = iter(obj.__pretty_repr__()) + header = next(iterator, None) + if not isinstance(header, PrettyType): + raise TypeError(f'First item must be PrettyType, got {type(header).__name__}') + if header.ref is not None: + _close(iterator) + return _Leaf(header.ref, 'ref') + type_name = header.type if isinstance(header.type, str) else header.type.__name__ + # Only named objects count as a nesting level; anonymous wrappers such + # as ``PrettyList`` / ``PrettyDict`` render like plain containers. + if type_name: + depth += 1 + max_depth = self.options.max_depth + if max_depth is not None and depth > max_depth + 1: + _close(iterator) + return _Leaf(f'{type_name}{header.start}...{header.end}') + + self.seen[key] = (obj, path) + self.active[key] = path + try: + group = _Group(type_name, header.start, header.end, [], header.empty_repr, len(header.elem_indent)) + attrs = [] + for elem in iterator: + if not isinstance(elem, PrettyAttr): + raise TypeError(f'Item must be Elem, got {type(elem).__name__}') + attrs.append(elem) + return self._build_items(group, attrs, path, depth, header.value_sep) + finally: + del self.active[key] + + +class _Layout: + """Lays out a document tree within the target width.""" + + def __init__(self, options: ReprOptions, color: bool): + self.options = options + self.color = color + self._widths: dict[int, float] = {} + + def flat_width(self, doc) -> float: + """Width of ``doc`` printed on one line (``inf`` if it cannot be).""" + key = id(doc) + width = self._widths.get(key) + if width is not None: + return width + if isinstance(doc, _Leaf): + width = math.inf if '\n' in doc.text else len(doc.text) + elif '\n' in doc.open or '\n' in doc.close: + width = math.inf + else: + width = len(doc.head) + len(doc.open) + len(doc.close) + if doc.items: + width += 2 * (len(doc.items) - 1) + for item in doc.items: + width += item.prefix_width() + self.flat_width(item.value) + len(item.end) + width += doc.single_comma + else: + width += len(doc.empty) + self._widths[key] = width + return width + + def _prefix(self, item: _Item) -> str: + if not item.key: + return item.start + return item.start + _style(item.key, 'key', self.color) + item.sep + + def flat(self, doc) -> str: + if isinstance(doc, _Leaf): + return _style(doc.text, doc.kind, self.color) + head = _style(doc.head, 'type', self.color) + if not doc.items: + return f'{head}{doc.open}{doc.empty}{doc.close}' + body = ', '.join(self._prefix(it) + self.flat(it.value) + it.end for it in doc.items) + if doc.single_comma: + body += ',' + return f'{head}{doc.open}{body}{doc.close}' + + def render(self, doc, indent: int, avail: float) -> str: + """Render ``doc`` at ``indent`` with ``avail`` columns left on its first line.""" + if isinstance(doc, _Leaf): + return _style(doc.text, doc.kind, self.color).replace('\n', '\n' + ' ' * indent) + if not doc.items or self.flat_width(doc) <= avail: + return self.flat(doc) + + inner = indent + doc.indent + pad = ' ' * inner + lines = [_style(doc.head, 'type', self.color) + doc.open] + for item in doc.items: + used = inner + len(item.label) + item.prefix_width() + len(item.end) + 1 + value = self.render(item.value, inner, self.options.width - used) + line = f'{pad}{item.label}{self._prefix(item)}{value}{item.end},' + comment = item.comment() if callable(item.comment) else item.comment + if comment: + line += _style(f' # {comment}', 'comment', self.color) + lines.append(line) + lines.append(' ' * indent + doc.close) + return '\n'.join(lines) class MappingReprMixin(Mapping[A, B]): @@ -261,10 +1034,7 @@ class MappingReprMixin(Mapping[A, B]): ... return pretty_repr_object(self) ... >>> print(PrettyDict({'a': 1, 'b': 2})) - { - 'a': 1, - 'b': 2 - } + {'a': 1, 'b': 2} """ def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]: @@ -300,10 +1070,7 @@ class PrettyMapping(PrettyRepr): >>> m = PrettyMapping({'a': 1, 'b': 2}, type_name='MyDict') >>> print(m) - MyDict{ - 'a': 1, - 'b': 2 - } + MyDict{'a': 1, 'b': 2} """ mapping: Mapping type_name: str = '' @@ -340,6 +1107,8 @@ class PrettyReprContext(threading.local): representing the same object multiple times. """ seen_modules_repr: dict[int, Any] | None = None + options: Optional['ReprOptions'] = None + render_depth: int = 0 CONTEXT = PrettyReprContext() @@ -438,8 +1207,9 @@ def yield_unique_pretty_repr_items( >>> node1.next = node2 ... >>> # This will handle circular reference gracefully - >>> for item in yield_unique_pretty_repr_items(node1): - ... print(item) + >>> items = list(yield_unique_pretty_repr_items(node1)) + >>> [type(item).__name__ for item in items] + ['PrettyType', 'PrettyAttr', 'PrettyAttr'] """ if repr_object is None: repr_object = _default_repr_object diff --git a/brainstate/util/_pretty_repr_layout_test.py b/brainstate/util/_pretty_repr_layout_test.py new file mode 100644 index 00000000..790d74b6 --- /dev/null +++ b/brainstate/util/_pretty_repr_layout_test.py @@ -0,0 +1,509 @@ +# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Tests for the width-aware layout, options, and formatting of the pretty-repr engine.""" + +import dataclasses +import os +import threading +import unittest +from unittest import mock + +import brainunit as u +import jax +import jax.numpy as jnp +import numpy as np + +from brainstate.util import ( + PrettyAttr, + PrettyRepr, + PrettyType, + ReprOptions, + ReprSummary, + format_array_like, + format_bytes, + get_repr_options, + pretty_repr_object, + pretty_format, + PrettyDict, + PrettyList, + repr_options, + set_repr_options, +) +from brainstate.util._pretty_repr import CONTEXT + + +class Obj(PrettyRepr): + """A configurable PrettyRepr used to drive the renderer.""" + + def __init__(self, name='Obj', **attrs): + self._name = name + self._attrs = attrs + + def __pretty_repr__(self): + yield PrettyType(self._name) + for k, v in self._attrs.items(): + yield PrettyAttr(k, v) + + +class Summarised(Obj): + """A PrettyRepr exposing a summary hook, like ``graph.Node``.""" + + def __init__(self, summary, **attrs): + super().__init__('Summarised', **attrs) + self._summary = summary + + def __pretty_repr_summary__(self): + return self._summary + + +class TestFlatVersusBroken(unittest.TestCase): + + def test_fits_on_one_line(self): + self.assertEqual(repr(Obj(a=1, b=[1, 2])), 'Obj(a=1, b=[1, 2])') + + def test_width_boundary(self): + obj = Obj(a=1, b=2) # 'Obj(a=1, b=2)' is 13 columns + with repr_options(width=13): + self.assertEqual(repr(obj), 'Obj(a=1, b=2)') + with repr_options(width=12): + self.assertEqual(repr(obj), 'Obj(\n a=1,\n b=2,\n)') + + def test_nested_only_outer_breaks(self): + obj = Obj(inner=Obj('In', x=1), long='x' * 30) + with repr_options(width=40): + self.assertEqual( + repr(obj), + 'Obj(\n inner=In(x=1),\n long=' + 'x' * 30 + ',\n)', + ) + + def test_nested_inner_breaks_too(self): + obj = Obj(inner=Obj('In', a='a' * 20, b='b' * 20)) + with repr_options(width=30): + self.assertEqual( + repr(obj), + 'Obj(\n inner=In(\n a=' + 'a' * 20 + ',\n b=' + 'b' * 20 + ',\n ),\n)', + ) + + def test_empty_object(self): + self.assertEqual(repr(Obj()), 'Obj()') + + def test_multiline_leaf_forces_break_and_reindents(self): + class Multi: + def __repr__(self): + return 'line1\nline2' + + self.assertEqual( + repr(Obj(m=Multi(), k=1)), + 'Obj(\n m=line1\n line2,\n k=1,\n)', + ) + + def test_positional_items(self): + class Pos(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('Pos') + yield PrettyAttr('', '(3,) -> (4,)') + yield PrettyAttr('k', 1) + + self.assertEqual(repr(Pos()), 'Pos((3,) -> (4,), k=1)') + + def test_attr_start_end_are_kept(self): + class Wrapped(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('W') + yield PrettyAttr('k', 1, start='<', end='>') + + self.assertEqual(repr(Wrapped()), 'W()') + + def test_explicit_comment_only_when_broken(self): + class Commented(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('C') + yield PrettyAttr('k', 1, comment='note') + + self.assertEqual(repr(Commented()), 'C(k=1)') + with repr_options(width=4): + self.assertEqual(repr(Commented()), 'C(\n k=1, # note\n)') + + def test_plain_containers(self): + obj = Obj(t=(1,), t2=(1, 2), l=[], d={'a': 'x'}, s='str') + self.assertEqual(repr(obj), "Obj(t=(1,), t2=(1, 2), l=[], d={'a': 'x'}, s=str)") + + def test_single_tuple_broken_has_no_double_comma(self): + obj = Obj(t=('a' * 30,)) + with repr_options(width=20): + self.assertEqual(repr(obj), "Obj(\n t=(\n '" + 'a' * 30 + "',\n ),\n)") + + def test_types_render_by_name(self): + self.assertEqual(repr(Obj(t=int, d=np.float32)), 'Obj(t=int, d=float32)') + + def test_generated_dataclass_repr_is_laid_out(self): + @dataclasses.dataclass + class DC: + a: int + b: str + hidden: int = dataclasses.field(default=0, repr=False) + + self.assertEqual(repr(Obj(dc=DC(1, 'x'))), "Obj(dc=DC(a=1, b='x'))") + + def test_handwritten_dataclass_repr_is_respected(self): + @dataclasses.dataclass + class DC: + a: int + + def __repr__(self): + return 'custom' + + self.assertEqual(repr(Obj(dc=DC(1))), 'Obj(dc=custom)') + + def test_invalid_item_raises(self): + class Bad(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('Bad') + yield 'oops' + + with self.assertRaises(TypeError): + repr(Bad()) + + +class TestPrettyFormat(unittest.TestCase): + + def test_plain_values(self): + self.assertEqual(pretty_format([np.zeros(2, np.float32), 'a', {'k': (1,)}]), "[f32[2], 'a', {'k': (1,)}]") + self.assertEqual(pretty_format(3), '3') + + def test_pretty_objects_and_width(self): + self.assertEqual(pretty_format(Obj(a=1)), 'Obj(a=1)') + with repr_options(width=6): + self.assertEqual(pretty_format([10, 20]), '[\n 10,\n 20,\n]') + + def test_color(self): + self.assertIn('\x1b[', pretty_format(Obj(a=1), color=True)) + + +class TestReferences(unittest.TestCase): + + def test_shared_reference(self): + child = Obj('Child', x=1) + self.assertEqual(repr(Obj(a=child, b=child)), 'Obj(a=Child(x=1), b=)') + + def test_shared_in_list_and_dict_paths(self): + child = Obj('Child') + self.assertEqual( + repr(Obj(l=[child], d={'k': child})), + "Obj(l=[Child()], d={'k': })", + ) + + def test_cycle_to_root(self): + obj = Obj(x=1) + obj._attrs['me'] = obj + self.assertEqual(repr(obj), 'Obj(x=1, me=)') + + def test_cycle_to_inner(self): + inner = Obj('In') + inner._attrs['parent'] = inner + self.assertEqual(repr(Obj(a=inner)), 'Obj(a=In(parent=))') + + def test_self_referencing_plain_containers(self): + lst = [1] + lst.append(lst) + dct = {'a': 1} + dct['self'] = dct + self.assertEqual(repr(Obj(l=lst, d=dct)), "Obj(l=[1, ], d={'a': 1, 'self': })") + + def test_self_referencing_dataclass(self): + @dataclasses.dataclass + class Link: + nxt: object = None + + link = Link() + link.nxt = link + self.assertEqual(repr(Obj(x=link)), 'Obj(x=Link(nxt=))') + + def test_repeated_plain_container_is_not_shared(self): + shared = [1, 2] + self.assertEqual(repr(Obj(a=shared, b=shared)), 'Obj(a=[1, 2], b=[1, 2])') + + def test_pretty_repr_returning_plain_iterable(self): + class ListRepr(PrettyRepr): + def __pretty_repr__(self): + return [PrettyType('L'), PrettyAttr('k', 1)] + + class EmptyRepr(PrettyRepr): + def __pretty_repr__(self): + return iter(()) + + self.assertEqual(repr(ListRepr()), 'L(k=1)') + with repr_options(max_depth=0): + self.assertEqual(repr(Obj(a=ListRepr())), 'Obj(a=L(...))') + with self.assertRaises(TypeError): + repr(EmptyRepr()) + + def test_header_ref(self): + class Ref(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('Ref', ref='') + + self.assertEqual(repr(Obj(r=Ref())), 'Obj(r=)') + + def test_sequence_of_objects_gets_index_labels_when_broken(self): + obj = Obj(layers=[Obj('A', v='a' * 30), Obj('B')]) + self.assertEqual(repr(obj), "Obj(layers=[A(v=" + 'a' * 30 + "), B()])") + with repr_options(width=30): + self.assertEqual( + repr(obj), + 'Obj(\n layers=[\n [0] A(\n v=' + 'a' * 30 + ',\n ),\n [1] B(),\n ],\n)', + ) + + +class TestLimits(unittest.TestCase): + + def test_max_depth(self): + obj = Obj(a=Obj('A', b=Obj('B', c=1)), l=[Obj('C')]) + with repr_options(max_depth=0): + self.assertEqual(repr(obj), 'Obj(a=A(...), l=[C(...)])') + with repr_options(max_depth=1): + self.assertEqual(repr(obj), 'Obj(a=A(b=B(...)), l=[C()])') + with repr_options(max_depth=None): + self.assertEqual(repr(obj), 'Obj(a=A(b=B(c=1)), l=[C()])') + + def test_max_items_for_objects_and_containers(self): + obj = Obj(**{f'k{i}': i for i in range(5)}) + with repr_options(max_items=2): + self.assertEqual(repr(obj), 'Obj(k0=0, k1=1, ... (+3 more))') + self.assertEqual(repr(Obj(l=list(range(5)))), 'Obj(l=[0, 1, ... (+3 more)])') + self.assertEqual(repr(Obj(d={1: 1, 2: 2, 3: 3})), 'Obj(d={1: 1, 2: 2, ... (+1 more)})') + + +class TestOptions(unittest.TestCase): + + def test_defaults(self): + options = ReprOptions() + self.assertEqual(options.width, 88) + self.assertIsNone(options.max_depth) + self.assertEqual(options.max_items, 64) + self.assertTrue(options.summary) + self.assertTrue(options.short_dtypes) + + def test_validation(self): + for kwargs in ( + dict(width=0), dict(width='88'), dict(max_depth=-1), + dict(max_items=0), dict(max_items=1.5), + ): + with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): + ReprOptions(**kwargs) + with self.assertRaises(TypeError): + with repr_options(unknown=1): + pass + + def test_context_nesting_and_restore_on_error(self): + base = get_repr_options() + with repr_options(width=50): + with repr_options(max_depth=2): + self.assertEqual(get_repr_options().width, 50) + self.assertEqual(get_repr_options().max_depth, 2) + self.assertIsNone(get_repr_options().max_depth) + with self.assertRaises(RuntimeError): + with repr_options(width=10): + raise RuntimeError + self.assertEqual(get_repr_options().width, 50) + self.assertEqual(get_repr_options(), base) + + def test_set_repr_options_global_and_restore(self): + previous = set_repr_options(width=33) + try: + self.assertEqual(get_repr_options().width, 33) + with repr_options(width=44): + self.assertEqual(get_repr_options().width, 44) + self.assertEqual(get_repr_options().width, 33) + finally: + set_repr_options(**dataclasses.asdict(previous)) + self.assertEqual(get_repr_options(), previous) + + def test_context_is_thread_local(self): + seen = {} + entered = threading.Event() + release = threading.Event() + + def other(): + entered.wait() + seen['width'] = get_repr_options().width + release.set() + + thread = threading.Thread(target=other) + thread.start() + with repr_options(width=21): + entered.set() + release.wait(5) + thread.join(5) + self.assertEqual(seen['width'], get_repr_options().width) + self.assertNotEqual(seen['width'], 21) + + def test_render_depth_restored_after_error(self): + class Boom(PrettyRepr): + def __pretty_repr__(self): + yield PrettyType('Boom') + raise RuntimeError('boom') + + with self.assertRaises(RuntimeError): + repr(Boom()) + self.assertEqual(CONTEXT.render_depth, 0) + + +class TestColor(unittest.TestCase): + + def test_repr_never_colored(self): + with repr_options(color=True): + self.assertNotIn('\x1b[', repr(Obj(a=np.zeros(2)))) + + def test_str_colored_when_enabled(self): + obj = Obj(a=np.zeros(2, dtype=np.float32), b=[Obj('C')]) + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop('NO_COLOR', None) + with repr_options(color=True): + text = str(obj) + self.assertIn('\x1b[1;36mObj\x1b[0m', text) + self.assertIn('\x1b[34ma\x1b[0m=', text) + self.assertIn('\x1b[32mf32[2]\x1b[0m', text) + + def test_str_plain_by_default_and_with_no_color(self): + obj = Obj(a=1) + self.assertEqual(str(obj), repr(obj)) + with mock.patch.dict(os.environ, {'NO_COLOR': '1'}): + with repr_options(color=True): + self.assertEqual(str(obj), 'Obj(a=1)') + + def test_str_respects_overridden_repr(self): + class Custom(Obj): + def __repr__(self): + return 'custom' + + with repr_options(color=True): + self.assertEqual(str(Custom()), 'custom') + + def test_str_of_containers_uses_engine(self): + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop('NO_COLOR', None) + with repr_options(color=True): + self.assertIn('\x1b[', str(PrettyList([Obj('C')]))) + self.assertIn('\x1b[', str(PrettyDict({'a': Obj('C')}))) + + def test_colored_layout_uses_plain_widths(self): + obj = Obj(a=1, b=2) + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop('NO_COLOR', None) + with repr_options(color=True, width=13): + text = str(obj) + self.assertNotIn('\n', text) # escape codes do not count toward width + + +class TestSummary(unittest.TestCase): + + def test_root_footer(self): + obj = Summarised(ReprSummary(26, 104), x=1) + self.assertEqual(repr(obj), 'Summarised(x=1) # 26 params (104 B)') + with repr_options(summary=False): + self.assertEqual(repr(obj), 'Summarised(x=1)') + + def test_footer_with_state_values_and_empty_summary(self): + obj = Summarised(ReprSummary(1, 4, 2, 8)) + self.assertEqual(repr(obj), 'Summarised() # 1 param (4 B), 2 state values (8 B)') + self.assertEqual(repr(Summarised(ReprSummary())), 'Summarised()') + + def test_child_comment_only_when_parent_broken(self): + obj = Obj(a=Summarised(ReprSummary(16, 64), w='w' * 20), b=Summarised(ReprSummary(0, 0, 3, 12))) + self.assertNotIn('#', repr(obj)) + with repr_options(width=40): + self.assertEqual( + repr(obj), + 'Obj(\n a=Summarised(w=' + 'w' * 20 + '), # 16 params\n' + ' b=Summarised(), # 3 state values\n)', + ) + + def test_singular_wording(self): + self.assertEqual(ReprSummary(1, 4).comment(), '1 param') + self.assertEqual(ReprSummary(0, 0, 1, 4).footer(), '1 state value (4 B)') + self.assertEqual(ReprSummary(1234, 4936).footer(), '1,234 params (4.8 KiB)') + + def test_failing_summary_is_ignored(self): + class Broken(Obj): + def __pretty_repr_summary__(self): + raise RuntimeError + + self.assertEqual(repr(Broken('B')), 'B()') + + def test_format_bytes(self): + self.assertEqual(format_bytes(0), '0 B') + self.assertEqual(format_bytes(1023), '1023 B') + self.assertEqual(format_bytes(1024), '1.0 KiB') + self.assertEqual(format_bytes(5 * 1024 ** 2), '5.0 MiB') + self.assertEqual(format_bytes(2048 * 1024 ** 3), '2048.0 GiB') + + +class TestFormatArrayLike(unittest.TestCase): + + def test_dtype_short_names(self): + cases = { + np.float16: 'f16', np.float32: 'f32', np.float64: 'f64', + np.int8: 'i8', np.int32: 'i32', np.uint8: 'u8', + np.complex64: 'c64', np.bool_: 'bool', + } + for dtype, name in cases.items(): + with self.subTest(dtype=dtype): + self.assertEqual(format_array_like(np.zeros((2, 3), dtype=dtype)), f'{name}[2,3]') + self.assertEqual(format_array_like(jnp.zeros(2, dtype=jnp.bfloat16)), 'bf16[2]') + + def test_long_dtype_names(self): + self.assertEqual(format_array_like(np.zeros(3, np.float32), short_dtypes=False), 'float32[3]') + with repr_options(short_dtypes=False): + self.assertEqual(repr(Obj(a=np.zeros(3, np.int32))), 'Obj(a=int32[3])') + + def test_scalars_and_abstract_values(self): + self.assertEqual(format_array_like(jnp.asarray(0.5, dtype=jnp.float32)), 'f32(0.5)') + self.assertEqual(format_array_like(np.bool_(True)), 'bool(True)') + self.assertEqual(format_array_like(jax.ShapeDtypeStruct((), jnp.float32)), 'f32[]') + self.assertEqual(format_array_like(jax.ShapeDtypeStruct((4, 5), jnp.int32)), 'i32[4,5]') + aval = jax.api_util.shaped_abstractify(jnp.zeros((2,))) + self.assertEqual(format_array_like(aval), 'f32[2]') + + def test_tracers_are_summarised(self): + captured = [] + + @jax.jit + def f(x): + captured.append(format_array_like(x)) + return x + + f(jnp.zeros((2, 2))) + self.assertEqual(captured, ['f32[2,2]']) + + def test_quantities(self): + self.assertEqual(format_array_like(jnp.ones(3) * u.mV), 'f32[3] mV') + self.assertEqual(format_array_like(u.Quantity(jnp.ones(2))), 'f32[2]') + self.assertEqual(format_array_like(5. * u.ms), '5. ms') + self.assertEqual(format_array_like(jnp.asarray(2.0) * u.mV), '2. mV') + + def test_typed_prng_keys(self): + self.assertEqual(format_array_like(jax.random.key(0)), 'key[]') + self.assertEqual(format_array_like(jax.random.split(jax.random.key(0), 3)), 'key[3]') + + def test_non_arrays(self): + for value in (1, 1.5, 'x', [1], None, object()): + with self.subTest(value=value): + self.assertIsNone(format_array_like(value)) + + +if __name__ == '__main__': + unittest.main() diff --git a/brainstate/util/_pretty_repr_test.py b/brainstate/util/_pretty_repr_test.py index 6aba0adc..ada3c46f 100644 --- a/brainstate/util/_pretty_repr_test.py +++ b/brainstate/util/_pretty_repr_test.py @@ -32,6 +32,7 @@ yield_unique_pretty_repr_items, _default_repr_object, _default_repr_attr, + repr_options, ) @@ -357,7 +358,9 @@ def __repr__(self): return pretty_repr_object(self) text = repr(PrettyDict({'a': 1, 'b': 2})) - self.assertEqual(text, "{\n 'a': 1,\n 'b': 2\n}") + self.assertEqual(text, "{'a': 1, 'b': 2}") + with repr_options(width=8): + self.assertEqual(repr(PrettyDict({'a': 1, 'b': 2})), "{\n 'a': 1,\n 'b': 2,\n}") class TestPrettyMapping(unittest.TestCase): diff --git a/docs/specs/pretty-repr-redesign.md b/docs/specs/pretty-repr-redesign.md new file mode 100644 index 00000000..331886f2 --- /dev/null +++ b/docs/specs/pretty-repr-redesign.md @@ -0,0 +1,131 @@ +# Pretty-repr redesign (Node / Module / State) + +## Goal + +Replace the one-attribute-per-line Flax-style renderer behind `graph.Node`, +`nn.Module`, `State`, `TreefyState` and the `brainstate.util` containers with a +width-aware compact tree renderer, plus a summary footer, configurable +options, and opt-in ANSI color. + +## Target output (width 88) + +``` +Sequential( + (3,) -> (2,), + layers=[ + [0] Linear((3,) -> (4,), weight=ParamState{bias: f32[4], weight: f32[3,4]}), # 16 params + [1] ReLU(), + [2] Linear((4,) -> (2,), weight=ParamState{bias: f32[2], weight: f32[4,2]}), # 10 params + ], +) # 26 params (104 B) +``` + +## Protocol (backward compatible) + +`__pretty_repr__` still yields one `PrettyType` followed by `PrettyAttr`s. +Additive fields: + +| field | meaning | +|-------------------------|-----------------------------------------------------------------| +| `PrettyAttr.key == ''` | positional item, rendered without `key` | +| `PrettyAttr.comment` | trailing `# comment`, shown only when the parent group is broken | +| `PrettyType.ref` | when set, the object renders as this literal marker | + +`PrettyAttr.value` that is a `str` is still inserted verbatim (protocol +unchanged); `PrettyObject` attribute walking now passes `repr(str)` so string +attributes are quoted (`name='w'`). + +## Rendering + +1. **Build** a document tree. `PrettyRepr` values, exact `list`, `tuple` and + `dict`, and dataclasses using the *generated* `__repr__` are expanded + natively; classes render by `__name__`; array-likes are summarised by + `format_array_like`; everything else is a `repr()` leaf (multi-line leaves + force a break). +2. **Layout** (black/Wadler rule): a group is printed flat when its flat width + fits in the remaining width and nothing inside forces a break. Otherwise it + is broken one item per line, indented by `len(PrettyType.elem_indent)`, + with trailing commas. Comments are emitted only on broken lines and never + count toward width. +3. **Lists/tuples** that contain `PrettyRepr` items get `[i] ` index labels + when broken. + +### Shared and cyclic references + +The renderer records the path (`.a`, `[0]`, `['k']`) of each `PrettyRepr` +object it expands. A repeat renders as ``; a reference to an +ancestor still being rendered renders as `` (`.` is the root). +Plain containers and dataclasses get cycle protection only (a repeated plain +list is printed again, not as ``). + +### Array formatting + +`jax.Array`, tracers, `np.ndarray`, `ShapedArray`, `ShapeDtypeStruct` render as +`f32[3,4]` (`bool`, `i32`, `u8`, `bf16`, `c64`, ...). 0-d concrete arrays +render as `f32(0.5)`. `brainunit.Quantity` appends its unit: `f32[3] mV`; 0-d +concrete quantities use `str(q)`. With `short_dtypes=False` the long dtype +name (`float32[3,4]`) is used. + +## Summary + +`Node.__pretty_repr_summary__()` returns a `ReprSummary(n_params, +param_bytes, n_state_values, state_bytes)` computed over the unique `State`s +reachable from the node (`ParamState` → params, other states → state +values). The renderer shows: + +- root footer: `) # 26 params (104 B)`; `, 10 state values (40 B)` appended + when non-param states exist; omitted when everything is zero; +- child comment on broken parents: `# 16 params` (or `# 10 state values`). + +Disabled with `summary=False`. + +## Options + +`ReprOptions(width=88, max_depth=None, max_items=64, summary=True, +color=False, short_dtypes=True)`. + +- `set_repr_options(**kw)` — process-wide default. +- `repr_options(**kw)` — thread-local context manager, nestable, restored on + exit/exception. +- `max_depth`: number of named-object levels expanded below the root; deeper + objects render as `Type(...)` (`0` collapses all children). Anonymous + wrappers (`PrettyList`, `PrettyDict`) do not count as a level. +- `max_items`: longer groups are cut with a `... (+N more)` item. +- `color`: ANSI styling, applied only by `str()` / `print()`, never by + `repr()`. Disabled whenever `NO_COLOR` is set. `BRAINSTATE_REPR_COLOR=1` + turns the default on at import. + +## Class-specific formats + +- `Module`: `in_size`/`out_size` collapse into one leading positional + `(3,) -> (4,)` when both are set. +- `State`: value is positional: `ParamState(f32[3])`, + `ParamState(f32[3], name='w')`; a dict value without other metadata uses + the brace form `ParamState{bias: f32[4], weight: f32[3,4]}`. +- `TreefyState`: value abstracted like `State`; `_hooks_manager` hidden; + `None` metadata hidden. +- `Map`/`Param` item hooks delegate to `Node.__pretty_repr_item__` first. + +## Formerly hand-written reprs + +- `nn/_elementwise.py` activations: f-string `__repr__`s removed; the Module + engine renders them (`LeakyReLU(negative_slope=0.01)`). +- `nn.Transform` and `interop.Transform`: `PrettyRepr` subclasses whose + subclasses report constructor arguments through `_repr_items()` (`(name, + value)` pairs, `''` = positional). Fixes `PermuteAxes` printing as + `Transpose`. +- `HiData`: `PrettyRepr`; prints the real subclass name; long non-structured + values are still truncated at 60 characters. +- `RandomState`: `RandomState(key[], data=[0, 42])`, read without hooks. +- `mixin.Mode` family: `Training`, `JointMode(Training, Batching(...))`; + `Batching` now labels its fields `batch_size`/`batch_axis` (was a + mislabelled `in_size`). `ParamDescriber` keeps call syntax but formats its + arguments with `pretty_format`. +- `JointTypes[...]` / `OneOfTypes[...]` keep typing-alias syntax. +- `PrettyRepr.__str__` only uses the engine when `__repr__` is engine-backed + (`__repr__._uses_pretty_repr = True`), so subclasses with their own + `__repr__` print consistently. + +## Out of scope + +HTML/treescope rendering, re-executing notebooks.