[chaoming0625]

Summary

Replaces the one-attribute-per-line repr engine with a width-aware compact tree layout shared by graph.Node, nn.Module, State, TreefyState, the brainstate.util containers, and the classes that previously had hand-written __repr__s.

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)

• Layout: one line when it fits (default width 88), otherwise one item per line with trailing commas. Sequences of objects get [i] labels, and repeated or self-referencing objects show as <shared: .a> / <cycle: .>.
• Arrays: f32[3,4], f32(0.5), f32[3] mV, key<fry>[].
• Summary: a root footer and per-child param/state counts. Shared states are counted once, and no state read hooks are triggered.
• Options: repr_options() (thread-local), set_repr_options(), get_repr_options(), and pretty_format() for arbitrary values. ANSI color is opt-in and applies to print() only, never repr(). NO_COLOR turns it off; BRAINSTATE_REPR_COLOR=1 turns it on by default.
• Engine-rendered now: elementwise activations, the nn and interop transforms (via a _repr_items() hook), HiData, RandomState (RandomState(key<fry>[], data=[0, 42])), the Mode family, and the arguments shown by ParamDescriber.
• Compatibility: the existing __pretty_repr__ protocol is unchanged. Additions are optional: an empty key means a positional item, plus PrettyAttr.comment and PrettyType.ref.

Bug fixes

• TreefyState printed full arrays and _hooks_manager=<object at 0x...>.
• Map / Param showed None-valued attributes.
• Batching labelled batch_size as in_size.
• PermuteAxes printed as Transpose(...).
• HiData subclasses printed as HiData(...).

Test plan

• Full suite: 5517 passed, 53 skipped.
• New brainstate/util/_pretty_repr_layout_test.py (55 tests), plus golden tests for Node, Module and State.
• Each bug fix has a test that failed before the fix.
• Doctests pass for the touched modules. Exceptions, all failing for reasons unrelated to this PR: check_state_jax_tracer, check_state_value_tree and the HiData clone example.
• Coverage measured with coverage run (pytest --cov crashes Python in my environment): _pretty_repr.py 99%, graph/_node.py 95%, _pretty_pytree.py 87%.

Not included

• The stored outputs in the docs notebooks still show the old format and need re-running.

Spec: docs/specs/pretty-repr-redesign.md

Summary by Sourcery

Replace the legacy verbose representations with a compact, width-aware pretty-printing system across the framework and its utility containers.

New Features:

• Introduce a width-aware tree renderer for nodes, modules, states, containers, and other supported objects, including compact layouts, references, array summaries, and configurable formatting options.

Bug Fixes:

• Correct state, container, transform, mode, and hierarchical-data representations by hiding internal or empty attributes, preserving concrete type names, and using accurate field labels.

Enhancements:

• Add parameter and state size summaries to graph and module representations without triggering state read hooks.
• Unify representations of transforms, activations, modes, random states, hierarchical data, and parameter descriptors through the shared pretty-repr protocol.
• Support thread-local representation settings, depth and item limits, optional colorized printing, and formatting of arbitrary values.

Documentation:

• Add a specification documenting the redesigned pretty-repr protocol, layout behavior, summaries, options, and class-specific formats.

Tests:

• Add comprehensive layout, formatting, reference-handling, option, summary, and array-formatting coverage, along with golden tests for nodes, modules, and states.

[sourcery-ai]

Sorry @chaoming0625, you've used your own review budget of 250,000 diff characters for the last 7 days.

You can request another review in 6 days by commenting @sourcery-ai review. Upgrade to get a review now.

[sourcery-ai]

Reviewer's Guide

Introduces a shared, width-aware pretty-representation engine and applies it throughout nodes, modules, states, containers, and previously custom-repr classes, adding compact layouts, array formatting, reference detection, summaries, options, and color-aware printing while retaining protocol compatibility.

Sequence diagram for width-aware representation rendering

sequenceDiagram
    participant Caller
    participant Engine as pretty_repr_object
    participant Object as PrettyRepr object
    participant Builder as DocumentBuilder
    participant Layout
    Caller->>Engine: pretty_repr_object(value)
    Engine->>Object: __pretty_repr__()
    Object-->>Builder: PrettyType and PrettyAttr items
    Builder->>Builder: format_array_like(value)
    Builder->>Builder: detect shared and cyclic references
    Builder-->>Layout: document tree
    Layout->>Layout: compare flat layout with ReprOptions.width
    Layout-->>Engine: compact or multiline text
    Engine-->>Caller: representation

Loading

Flow diagram for pretty-repr rendering

flowchart TD
    Value["Value or PrettyRepr object"] --> Build["Build document tree"]
    Build --> Detect["Format arrays and detect shared/cyclic references"]
    Detect --> Measure["Measure against configured width"]
    Measure -->|Fits| Compact["Render compact layout"]
    Measure -->|Does not fit| Expanded["Render multiline layout with labels and commas"]
    Compact --> Summary["Append optional parameter/state summary"]
    Expanded --> Summary
    Summary --> Result["repr() or color-aware print()"]

Loading

File-Level Changes

Change	Details	Files
Replaced the legacy line-oriented renderer with a shared width-aware document/layout engine and configurable formatting API.	Added compact-vs-multiline layout with trailing commas, sequence labels, cycle/shared-reference markers, depth/item limits, summaries, and optional print-only ANSI color. Added array, scalar, quantity, typed-key, and byte-size formatting plus arbitrary-value support through pretty_format(). Added thread-local and process-wide representation options while preserving the existing pretty_repr protocol with additive metadata fields.	brainstate/util/_pretty_repr.py brainstate/util/_pretty_pytree.py brainstate/util/__init__.py brainstate/util/_pretty_repr_layout_test.py brainstate/util/_pretty_repr_test.py docs/specs/pretty-repr-redesign.md
Integrated the new renderer into graph nodes, modules, states, containers, and parameter/state summaries.	Changed State and TreefyState output to abstract array values, use positional/brace forms, hide internal metadata, and avoid triggering read hooks. Added Node-level unique parameter/state counting and root/child summary comments; Module merges in/out sizes into a positional shape item. Updated container repr handling for strings, nested state values, index labels, and shared object references; Map and Param now inherit filtering behavior.	brainstate/_state.py brainstate/_state_test.py brainstate/graph/_node.py brainstate/graph/_node_test.py brainstate/nn/_module.py brainstate/nn/_module_test.py brainstate/nn/_common.py brainstate/nn/_common_test.py brainstate/nn/_param.py brainstate/nn/_param_test.py brainstate/util/_pretty_pytree.py
Migrated formerly hand-written representations to engine-backed protocols and hooks across neural-network, interop, random, data, and mode classes.	Removed activation f-string reprs and added _repr_items() hooks for nn and interop transforms, fixing transform naming and positional argument rendering. Rendered HiData, RandomState, Mode/JointMode/Batching, and ParamDescriber through the shared engine with corrected field labels and raw-key display. Added regression coverage for the migrated representations and documented the supported class-specific formats.	brainstate/nn/_elementwise.py brainstate/nn/_hidata.py brainstate/nn/_hidata_test.py brainstate/nn/_transform.py brainstate/interop/_transforms.py brainstate/mixin.py brainstate/mixin_test.py brainstate/random/_state.py brainstate/random/_state_test.py brainstate/nn/_common.py docs/specs/pretty-repr-redesign.md

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[codecov]

Codecov Report

❌ Patch coverage is 89.34708% with 62 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
brainstate/util/_pretty_pytree.py	30.30%	46 Missing ⚠️
brainstate/nn/_transform.py	89.18%	4 Missing ⚠️
brainstate/util/_pretty_repr.py	98.87%	1 Missing and 3 partials ⚠️
brainstate/_state.py	90.00%	2 Missing and 1 partial ⚠️
brainstate/random/_state.py	70.00%	1 Missing and 2 partials ⚠️
brainstate/graph/_node.py	92.59%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!