Skip to content

Add rich HTML formatter for PyTorch nn.Module#8307

Open
manzt wants to merge 4 commits intomainfrom
push-uynvsxvzkrol
Open

Add rich HTML formatter for PyTorch nn.Module#8307
manzt wants to merge 4 commits intomainfrom
push-uynvsxvzkrol

Conversation

@manzt
Copy link
Contributor

@manzt manzt commented Feb 13, 2026

Adds a formatter that renders nn.Module instances as rich, collapsible HTML trees.

The formatter uses PyTorch's subpackage structure to map known modules to four display categories (Weight, Activation, Normalization, Regularization). Custom modules that follow PyTorch conventions work automatically, and unknown/containers default to "Other".

Keyword-heavy layers (Conv2d with 8+ args) get an expandable row that shows the full args truncated with CSS ellipsis, expandable to a line-by-line view. Positional args (like channel dimensions) stay always-visible next to the type pill for quick scanning. Frozen layers are visually dimmed and show "(frozen)" next to their param count, while partially-frozen containers show the trainable subset.

A hover legend in the footer explains the category color coding using mini swatch replicas of the actual type pills.

Hovering legend:

Adds a formatter that renders `nn.Module` instances as rich, collapsible
HTML trees.

The formatter uses PyTorch's subpackage structure to map known modules
to four display categories (Weight, Activation, Normalization,
Regularization). Custom modules that follow PyTorch conventions work
automatically, and unknown/containers default to "Other".

Keyword-heavy layers (Conv2d with 8+ args) get an expandable row that
shows the full args truncated with CSS ellipsis, expandable to a
line-by-line view. Positional args (like channel dimensions) stay
always-visible next to the type pill for quick scanning. Frozen layers
are visually dimmed and show "(frozen)" next to their param count, while
partially-frozen containers show the trainable subset.

A hover legend in the footer explains the category color coding using
mini swatch replicas of the actual type pills.
@vercel
Copy link

vercel bot commented Feb 13, 2026

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
marimo-docs Ready Ready Preview, Comment Feb 13, 2026 8:34pm

Request Review

@manzt manzt requested a review from koaning February 13, 2026 19:58
@dmadisetti
Copy link
Collaborator

😍

Checked in darkmode too:

image

@koaning
Copy link
Contributor

koaning commented Feb 13, 2026

Could we also show the number of trainable and non-trainable weights? This is relevant when you're keen to do some finetuning and you want to make sure that you're only taking a subset of all weights to train the bigger system.

Nevermind, found it!

CleanShot 2026-02-13 at 22 51 50@2x

Maybe the dtype of the layer too? Maybe what device (cpu/gpu) the element is on?

The general feeling is that it is really sweet :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants