tensorgami 0.1.0__py3-none-any.whl

tensorgami/__init__.py

@@ -0,0 +1,22 @@
+"""Tensorgami: draw the architecture and information flow of ML models."""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+
+def __getattr__(name: str):
+    """Lazy top-level API, so ``import tensorgami`` stays light: matplotlib loads when a
+    figure is made, torch when a model is drawn."""
+    if name in ("Figure", "Input", "Loss", "Note", "draw"):
+        from tensorgami import api
+
+        return getattr(api, name)
+    raise AttributeError(f"module 'tensorgami' has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:  # tab completion lists the lazy names too
+    return sorted({*globals(), *__all__})
+
+
+__all__ = ["Figure", "Input", "Loss", "Note", "__version__", "draw"]

tensorgami/api/__init__.py

@@ -0,0 +1,9 @@
+"""Public API: the objects users work with (ADR-0007)."""
+
+from __future__ import annotations
+
+from tensorgami.api.annotations import Input, Loss, Note
+from tensorgami.api.draw import draw
+from tensorgami.api.figure import Figure
+
+__all__ = ["Figure", "Input", "Loss", "Note", "draw"]

tensorgami/api/annotations.py

@@ -0,0 +1,43 @@
+"""Annotations: what a figure shows beyond the traced modules (ADR-0007 ``annotations=``)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, init=False)
+class Loss:
+    """A loss term: a box fed by the tensors it compares, e.g.
+    ``tg.Loss("reconstruction", "output 0", "input")``. ``trains`` names the models or modules
+    its gradient updates; dashed arrows go back to them (the adversarial loop of a GAN)."""
+
+    name: str
+    inputs: tuple[str, ...]
+    trains: tuple[str, ...]
+
+    def __init__(self, name: str, *inputs: str, trains: tuple[str, ...] | list[str] = ()) -> None:
+        if not inputs:
+            raise ValueError(f"Loss({name!r}) needs the tensors it compares, e.g. 'output'")
+        object.__setattr__(self, "name", name)
+        object.__setattr__(self, "inputs", tuple(inputs))
+        object.__setattr__(self, "trains", tuple(trains))
+
+
+@dataclass(frozen=True)
+class Input:
+    """Another input for a tensor of the figure: it feeds what that tensor feeds, e.g. the real
+    images a discriminator also judges, ``tg.Input("real image", to="discriminator.image")``."""
+
+    name: str
+    to: str
+
+
+@dataclass(frozen=True)
+class Note:
+    """A line of text under a node, e.g. ``tg.Note("ε ~ N(0, I)", at="...")``."""
+
+    text: str
+    at: str
+
+
+Annotation = Loss | Input | Note

tensorgami/api/draw.py

@@ -0,0 +1,187 @@
+"""``tg.draw``: from a loaded model to a figure in one call (ADR-0007, ADR-0008)."""
+
+from __future__ import annotations
+
+import difflib
+from collections.abc import Mapping, Sequence
+from dataclasses import fields
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from tensorgami.api.annotations import Annotation, Input, Loss, Note
+from tensorgami.api.figure import Figure
+from tensorgami.ir import Graph
+from tensorgami.layout import LayoutOptions
+from tensorgami.style import DEFAULT_STYLE, Theme, load_style
+from tensorgami.transforms import Level
+from tensorgami.transforms.annotate import add_input, add_loss, add_note, compose, link
+from tensorgami.views import get_view
+
+if TYPE_CHECKING:
+    from torch import nn
+
+LAYOUT_OPTIONS = tuple(f.name for f in fields(LayoutOptions))
+_PARAMETERS = (
+    "level",
+    "view",
+    "style",
+    "overrides",
+    "links",
+    "annotations",
+    "direction",
+    *LAYOUT_OPTIONS,
+)
+_RENAMED = {"theme": "style"}  # the name sketched in ADR-0007, before ADR-0008
+
+
+def draw(
+    model: nn.Module | Mapping[str, nn.Module],
+    example_input: Any,
+    *,
+    level: int | str = Level.LAYERS,
+    view: str = "auto",
+    style: str | Path | Theme = DEFAULT_STYLE,
+    overrides: Mapping[str, Mapping[str, Any]] | None = None,
+    links: Sequence[tuple[str, ...]] = (),
+    annotations: Sequence[Annotation] = (),
+    direction: str | None = None,
+    **layout_options: Any,
+) -> Figure:
+    """Draw ``model`` as it runs on ``example_input``.
+
+    The model runs once, without gradients and in eval mode; its weights, training flags and
+    the caller's random state are left as they were. Nothing is written to disk: the figure
+    shows inline in Jupyter and is saved only on request, ``fig.save("model.png", dpi=300)``.
+
+    Args:
+        model: a loaded ``torch.nn.Module``, on any device (the input is moved to it).
+        example_input: what the model is called with: a tensor, a tuple of positional
+            arguments, or a dict of keyword arguments. Include the batch dimension.
+        level: the detail, by number or name: 0 ``"system"`` (one box), 1 ``"stages"`` (the
+            model's direct sub-modules), 2 ``"layers"`` (default), 3 ``"operators"``
+            (everything traced), 4 ``"cells"``.
+        view: ``"block"`` (boxes and wires), ``"unrolled"`` (recurrence as a time × layer
+            grid, with per-step inputs x_t), ``"volume"`` (CNN layers as tensor volumes),
+            ``"cell"`` (the recurrent cells drawn as their gates), or ``"auto"`` (default:
+            unrolled when a module is called step after step with state passed along, block
+            otherwise).
+        style: a built-in style (``"default"``, ``"origami"``), the path of a ``.toml``
+            style, or a loaded :class:`~tensorgami.style.Theme`.
+        overrides: category, label or leaf-ness per module path or class name, e.g.
+            ``{"ConvLSTMCell": {"category": "recurrent", "label": "ConvLSTM cell"}}``.
+        links: with several models (``model`` a dict of name → module, ``example_input`` a
+            dict of name → input), which tensor feeds which model input, as ``(source,
+            target)`` or ``(source, target, label)``, e.g. ``("generator",
+            "discriminator.image")``: the generator's output is the image its discriminator
+            judges.
+        annotations: what no module contains: ``tg.Loss`` terms (with ``trains=`` for the
+            arrows back to the models they train), extra ``tg.Input`` and ``tg.Note``.
+            Tensors are named by input/output names (``"input"``, ``"output 0"``, keyword
+            names) or module paths (``"mu"``); with several models, after the model's name
+            (``"discriminator.image"``; the name alone is the model's output).
+        direction: ``"LR"`` (left → right) or ``"TB"`` (top → bottom); by default the
+            style's choice, else ``"LR"``.
+        **layout_options: further :class:`~tensorgami.layout.LayoutOptions`, e.g.
+            ``legend=False``, ``wrap_width=None`` (one row) or ``rank_gap=64``.
+
+    Example::
+
+        import tensorgami as tg
+
+        fig = tg.draw(model, torch.randn(1, 3, 224, 224), style="origami")
+        fig  # shows inline in Jupyter
+        fig.save("model.png", dpi=300)
+    """
+    # every argument is checked before the model runs, so a typo costs no forward pass
+    _check_names(layout_options)
+    _check_extras(model, example_input, links, annotations)
+    theme = style if isinstance(style, Theme) else load_style(style)
+    if direction is not None:
+        layout_options["direction"] = direction
+    options = LayoutOptions.for_style(theme, **layout_options)
+    level = Level.parse(level)
+    chosen = get_view(view)
+
+    graph = _annotate(_capture(model, example_input, overrides), links, annotations)
+    shown, scene = chosen.draw(graph, theme, level=level, options=options)
+    return Figure(scene, theme, shown)
+
+
+def _capture(model, example_input: Any, overrides: Mapping | None) -> Graph:
+    from tensorgami.importers.torch import trace  # torch is imported only to draw a model
+
+    if not isinstance(model, Mapping):
+        return trace(model, example_input, overrides=overrides)
+    graphs = {
+        name: trace(module, example_input[name], overrides=_own(name, model, overrides))
+        for name, module in model.items()
+    }
+    return compose(graphs)
+
+
+def _own(name: str, models: Mapping, overrides: Mapping | None) -> dict | None:
+    """The overrides for one of several models: its own (``"generator.fc"`` → ``"fc"``) and
+    the shared ones (class names and paths without a model's name)."""
+    if not overrides:
+        return None
+    own = {}
+    for key, value in overrides.items():
+        head, _, rest = key.partition(".")
+        if head == name and rest:
+            own[rest] = value
+        elif head not in models:
+            own[key] = value
+    return own
+
+
+def _annotate(graph: Graph, links: Sequence[tuple[str, ...]], annotations) -> Graph:
+    for a in annotations:  # extra inputs first: a link may then replace the input they join
+        if isinstance(a, Input):
+            graph = add_input(graph, a.name, a.to)
+    for source, target, *label in links:
+        graph = link(graph, source, target, *label)
+    for a in annotations:
+        if isinstance(a, Loss):
+            graph = add_loss(graph, a.name, a.inputs, a.trains)
+        elif isinstance(a, Note):
+            graph = add_note(graph, a.text, a.at)
+    return graph
+
+
+def _check_extras(model, example_input, links, annotations) -> None:
+    if isinstance(model, Mapping):
+        if not model or not all(isinstance(k, str) and "." not in k for k in model):
+            raise ValueError("several models: a dict of name → model, names without dots")
+        missing = (
+            sorted(set(model) - set(example_input)) if isinstance(example_input, Mapping) else []
+        )
+        if not isinstance(example_input, Mapping) or missing:
+            raise ValueError(
+                "several models need a dict of example inputs with the same names "
+                f"({', '.join(model)}); missing: {', '.join(missing) or 'all'}"
+            )
+    for entry in links:
+        if not (isinstance(entry, tuple) and len(entry) in (2, 3)):
+            raise TypeError(
+                f"links: each entry is (source, target) or (source, target, label); got {entry!r}"
+            )
+    for a in annotations:
+        if not isinstance(a, Loss | Input | Note):
+            raise TypeError(f"annotations: tg.Loss, tg.Input or tg.Note; got {a!r}")
+
+
+def _check_names(layout_options: Mapping[str, Any]) -> None:
+    for name in layout_options:
+        if name in LAYOUT_OPTIONS:
+            continue
+        close = (
+            [_RENAMED[name]]
+            if name in _RENAMED
+            else difflib.get_close_matches(name, _PARAMETERS, n=1)
+        )
+        hint = (
+            f"; did you mean {close[0]!r}?"
+            if close
+            else f"; layout options are: {', '.join(LAYOUT_OPTIONS)}"
+        )
+        raise TypeError(f"draw() got an unexpected keyword argument {name!r}{hint}")

tensorgami/api/figure.py

@@ -0,0 +1,119 @@
+"""``tg.Figure``: what the user gets back — shown inline in Jupyter, saved only on request
+(ADR-0005, ADR-0007)."""
+
+from __future__ import annotations
+
+import io
+from functools import cached_property
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from tensorgami.ir import Graph
+from tensorgami.render.mpl import PX_PER_INCH, draw_scene, quiet_fonts, render_figure
+from tensorgami.render.scene_renderer import render_svg
+from tensorgami.scene import Scene
+from tensorgami.style import Theme
+
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+    from matplotlib.figure import Figure as MplFigure
+
+#: formats written through matplotlib (SVG is written by Tensorgami's own SVG canvas)
+RASTER_AND_PDF = {".png", ".pdf", ".jpg", ".jpeg", ".tif", ".tiff"}
+SCREEN_DPI = 2 * PX_PER_INCH  # notebook display: rendered at 2× for sharp HiDPI screens
+
+
+class Figure:
+    """A drawn model. Nothing is written to disk until :meth:`save` is called.
+
+    * in Jupyter, the figure renders inline (``_repr_png_``);
+    * ``fig.save("model.png", dpi=300)`` — also ``.svg`` and ``.pdf``;
+    * ``fig.show()`` opens a window from a plain script;
+    * ``fig.figure`` / ``fig.ax`` are the underlying matplotlib objects (titles, annotations);
+    * ``fig.draw_on(ax)`` draws into a panel of your own ``plt.subplots`` layout.
+    """
+
+    def __init__(self, scene: Scene, style: Theme, graph: Graph | None = None) -> None:
+        self._scene = scene
+        self._style = style
+        self._graph = graph
+
+    # ------------------------------------------------------------------ inspection
+
+    @property
+    def scene(self) -> Scene:
+        return self._scene
+
+    @property
+    def style(self) -> Theme:
+        return self._style
+
+    @property
+    def graph(self) -> Graph | None:
+        return self._graph
+
+    @property
+    def size(self) -> tuple[float, float]:
+        """Natural size in px (1 px = 1/96 inch)."""
+        return self._scene.width, self._scene.height
+
+    @cached_property
+    def figure(self) -> MplFigure:
+        return render_figure(self._scene, self._style)
+
+    @property
+    def ax(self) -> Axes:
+        return self.figure.axes[0]
+
+    def __repr__(self) -> str:
+        w, h = self.size
+        title, style = self._scene.title, self._style.name
+        return f"<tensorgami.Figure {title!r} {w:.0f}×{h:.0f} px, style {style!r}>"
+
+    # ------------------------------------------------------------------ output
+
+    def save(self, path: str | Path, *, dpi: float = 300) -> Path:
+        """Write the figure; the format follows the suffix: ``.png`` (default choice for
+        reports), ``.svg`` (vector, editable), ``.pdf``, ``.jpg``, ``.tif``."""
+        path = Path(path)
+        suffix = path.suffix.lower()
+        if suffix == ".svg":
+            path.write_text(self.to_svg(), encoding="utf-8")
+        elif suffix in RASTER_AND_PDF:
+            with quiet_fonts():
+                self.figure.savefig(path, dpi=dpi)
+        else:
+            raise ValueError(
+                f"cannot save as {suffix or 'a file without suffix'}: use .png, .svg, .pdf, "
+                ".jpg or .tif"
+            )
+        return path
+
+    def to_png(self, *, dpi: float = SCREEN_DPI) -> bytes:
+        buffer = io.BytesIO()
+        with quiet_fonts():
+            self.figure.savefig(buffer, format="png", dpi=dpi)
+        return buffer.getvalue()
+
+    def to_svg(self) -> str:
+        return render_svg(self._scene, self._style)
+
+    def _repr_png_(self) -> tuple[bytes, dict[str, int]]:
+        """Jupyter display: rendered at 2× but shown at the natural size."""
+        w, h = self.size
+        return self.to_png(), {"width": round(w), "height": round(h)}
+
+    def draw_on(self, ax: Axes) -> Axes:
+        """Draw into an existing Axes (e.g. a panel of ``plt.subplots``); returns it."""
+        draw_scene(ax, self._scene, self._style)
+        return ax
+
+    def show(self) -> None:
+        """Open the figure in a matplotlib window (from scripts)."""
+        import matplotlib.pyplot as plt
+
+        w, h = self.size
+        window = plt.figure(figsize=(w / PX_PER_INCH, h / PX_PER_INCH))
+        self.draw_on(window.add_axes((0.0, 0.0, 1.0, 1.0)))
+        with quiet_fonts():
+            plt.show()

tensorgami/importers/__init__.py

@@ -0,0 +1,13 @@
+"""Model importers: framework → IR (ADR-0001).
+
+Each framework lives in its own subpackage and is imported on demand, so ``tensorgami`` itself
+works without any deep-learning framework installed::
+
+    from tensorgami.importers.torch import trace
+"""
+
+from __future__ import annotations
+
+from tensorgami.importers.base import ImporterError, ModelImporter, UnclassifiedOpWarning
+
+__all__ = ["ImporterError", "ModelImporter", "UnclassifiedOpWarning"]

tensorgami/importers/base.py

@@ -0,0 +1,21 @@
+"""Framework-independent importer contract."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from tensorgami.ir import Graph
+
+
+class ImporterError(RuntimeError):
+    """The model could not be captured (e.g. it failed on the example input)."""
+
+
+class UnclassifiedOpWarning(UserWarning):
+    """Some operations are not in the catalogue and were drawn with the 'other' category."""
+
+
+class ModelImporter(Protocol):
+    """Turns a framework model plus an example input into an :class:`~tensorgami.ir.Graph`."""
+
+    def import_model(self, model: Any, example_input: Any) -> Graph: ...

tensorgami/importers/torch/__init__.py

@@ -0,0 +1,15 @@
+"""PyTorch importer (ADR-0003): ``trace(model, example_input)`` → :class:`tensorgami.ir.Graph`."""
+
+from __future__ import annotations
+
+try:
+    import torch  # noqa: F401
+except ImportError as exc:  # pragma: no cover - exercised only without torch
+    raise ImportError(
+        "tensorgami.importers.torch needs PyTorch. Install it the way your project does "
+        "(e.g. `pip install torch`), or use the extra: `pip install tensorgami[torch]`."
+    ) from exc
+
+from tensorgami.importers.torch.importer import TorchImporter, trace
+
+__all__ = ["TorchImporter", "trace"]

tensorgami/importers/torch/catalog.py

@@ -0,0 +1,215 @@
+"""What Tensorgami knows about PyTorch operations: names, categories, leaves, parameters.
+
+Pure data plus small lookup functions. Classes are matched by *name* along the MRO, so
+subclasses (e.g. ``NonDynamicallyQuantizableLinear``) inherit their parent's category and the
+catalogue never needs the class objects themselves. Name lists are whitespace-separated
+strings to keep the tables compact.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from tensorgami.vocabulary import OpCategory as C
+
+
+def _names(text: str) -> tuple[str, ...]:
+    return tuple(text.split())
+
+
+# --------------------------------------------------------------------------- modules
+
+_MODULES: dict[C, tuple[str, ...]] = {
+    C.CONV: _names("""Conv1d Conv2d Conv3d ConvTranspose1d ConvTranspose2d ConvTranspose3d
+                      LazyConv1d LazyConv2d LazyConv3d"""),
+    C.LINEAR: _names("Linear LazyLinear Bilinear"),
+    C.RECURRENT: _names("RNN LSTM GRU RNNCell LSTMCell GRUCell"),
+    C.ATTENTION: _names("MultiheadAttention"),
+    C.NORM: _names("""BatchNorm1d BatchNorm2d BatchNorm3d SyncBatchNorm LayerNorm GroupNorm
+                      InstanceNorm1d InstanceNorm2d InstanceNorm3d LocalResponseNorm RMSNorm"""),
+    C.ACTIVATION: _names("""ReLU ReLU6 LeakyReLU PReLU RReLU ELU SELU CELU GELU SiLU Mish
+                            Sigmoid Tanh Softmax LogSoftmax Softplus Softsign Hardtanh
+                            Hardswish Hardsigmoid GLU Tanhshrink Threshold"""),
+    C.RESAMPLE: _names("""MaxPool1d MaxPool2d MaxPool3d AvgPool1d AvgPool2d AvgPool3d
+                          AdaptiveAvgPool1d AdaptiveAvgPool2d AdaptiveAvgPool3d
+                          AdaptiveMaxPool1d AdaptiveMaxPool2d AdaptiveMaxPool3d Upsample
+                          UpsamplingNearest2d UpsamplingBilinear2d PixelShuffle
+                          PixelUnshuffle"""),
+    C.EMBEDDING: _names("Embedding EmbeddingBag"),
+    C.REGULARIZATION: _names(
+        "Dropout Dropout1d Dropout2d Dropout3d AlphaDropout FeatureAlphaDropout"
+    ),
+    C.RESHAPE: _names("""Flatten Unflatten ZeroPad1d ZeroPad2d ZeroPad3d ConstantPad1d
+                         ConstantPad2d ConstantPad3d ReflectionPad1d ReflectionPad2d
+                         ReflectionPad3d ReplicationPad1d ReplicationPad2d ReplicationPad3d"""),
+    C.OTHER: _names("Identity"),
+}
+MODULE_CATEGORY: dict[str, C] = {name: cat for cat, names in _MODULES.items() for name in names}
+
+#: Modules drawn as one operation; their internals are not traced (ADR-0003).
+LEAF_MODULES: frozenset[str] = frozenset(MODULE_CATEGORY)
+
+#: Hyper-parameters worth keeping per module type (first match along the MRO).
+MODULE_PARAMS: dict[str, tuple[str, ...]] = {
+    "_ConvNd": _names("in_channels out_channels kernel_size stride padding dilation groups"),
+    "Linear": _names("in_features out_features"),
+    "Bilinear": _names("in1_features in2_features out_features"),
+    "RNNBase": _names("input_size hidden_size num_layers bidirectional batch_first"),
+    "RNNCellBase": _names("input_size hidden_size"),
+    "MultiheadAttention": _names("embed_dim num_heads batch_first"),
+    "_NormBase": _names("num_features"),
+    "LayerNorm": _names("normalized_shape"),
+    "GroupNorm": _names("num_groups num_channels"),
+    "_MaxPoolNd": _names("kernel_size stride"),
+    "_AvgPoolNd": _names("kernel_size stride"),
+    "_AdaptiveAvgPoolNd": _names("output_size"),
+    "_AdaptiveMaxPoolNd": _names("output_size"),
+    "Upsample": _names("scale_factor size mode"),
+    "Embedding": _names("num_embeddings embedding_dim"),
+    "EmbeddingBag": _names("num_embeddings embedding_dim"),
+    "_DropoutNd": _names("p"),
+    "LeakyReLU": _names("negative_slope"),
+    "Softmax": _names("dim"),
+    "LogSoftmax": _names("dim"),
+    "Flatten": _names("start_dim end_dim"),
+}
+
+# --------------------------------------------------------------------------- functions
+
+_FUNCTIONS: dict[C, tuple[str, ...]] = {
+    C.CONV: _names("""conv1d conv2d conv3d conv_transpose1d conv_transpose2d conv_transpose3d
+                      convolution"""),
+    C.LINEAR: _names("linear bilinear matmul mm bmm addmm baddbmm einsum mv addmv"),
+    C.ACTIVATION: _names("""relu relu6 leaky_relu elu selu celu gelu silu mish sigmoid tanh
+                            softmax log_softmax softplus softsign hardtanh hardswish hardsigmoid
+                            glu logsigmoid prelu rrelu threshold tanhshrink exp log"""),
+    C.NORM: _names("""batch_norm layer_norm group_norm instance_norm local_response_norm
+                      normalize rms_norm"""),
+    C.RESAMPLE: _names("""max_pool1d max_pool2d max_pool3d avg_pool1d avg_pool2d avg_pool3d
+                          adaptive_avg_pool1d adaptive_avg_pool2d adaptive_avg_pool3d
+                          adaptive_max_pool1d adaptive_max_pool2d adaptive_max_pool3d
+                          interpolate upsample upsample_nearest upsample_bilinear pixel_shuffle
+                          pixel_unshuffle grid_sample"""),
+    C.ATTENTION: _names("scaled_dot_product_attention multi_head_attention_forward"),
+    C.EMBEDDING: _names("embedding embedding_bag one_hot"),
+    C.REGULARIZATION: _names("""dropout dropout1d dropout2d dropout3d alpha_dropout
+                                feature_alpha_dropout"""),
+    C.RESHAPE: _names("""view reshape flatten unflatten permute transpose t squeeze unsqueeze
+                         contiguous getitem split chunk unbind expand expand_as repeat tile
+                         narrow select movedim swapaxes swapdims view_as reshape_as flip roll
+                         pad to float double half bfloat16 int long bool type clone detach T mT
+                         H mH index_select gather masked_select data real imag"""),
+    C.MERGE: _names("cat concat concatenate stack hstack vstack dstack"),
+    C.STOCHASTIC: _names("""randn rand randint randn_like rand_like randint_like normal
+                            bernoulli multinomial randperm poisson"""),
+    C.TENSOR: _names("""zeros ones full empty zeros_like ones_like full_like empty_like arange
+                        linspace logspace eye tensor as_tensor new_zeros new_ones new_full
+                        new_empty new_tensor tril triu"""),
+    C.LOSS: _names("""mse_loss l1_loss cross_entropy nll_loss binary_cross_entropy
+                      binary_cross_entropy_with_logits kl_div smooth_l1_loss huber_loss
+                      cosine_embedding_loss margin_ranking_loss triplet_margin_loss"""),
+}
+FUNCTION_CATEGORY: dict[str, C] = {n: cat for cat, names in _FUNCTIONS.items() for n in names}
+
+#: Element-wise binary operations: a *merge* when they combine two tensors, plain arithmetic
+#: (category 'other') when one operand is a Python number, e.g. ``0.5 * logvar``.
+ARITHMETIC: frozenset[str] = frozenset(
+    _names("""add sub mul div true_divide floor_divide maximum minimum where remainder fmod
+              atan2 hypot logaddexp""")
+)
+
+#: Known operations without a dedicated category: drawn as 'other', without a warning.
+KNOWN_OTHER: frozenset[str] = frozenset(
+    _names("""neg abs sqrt rsqrt pow square reciprocal mean sum prod max min amax amin argmax
+              argmin std var norm clamp clip sign round floor ceil trunc cumsum cumprod
+              logsumexp sort topk eq ne lt le gt ge logical_not logical_and logical_or
+              logical_xor masked_fill fill copy zero index_put scatter scatter_add nan_to_num
+              isnan isinf sin cos erf all any count_nonzero diff outer cross dot trace diag
+              diagonal setitem""")
+)
+
+#: Operations that read their tensor argument only for its *shape*: they become sources
+#: without data edges (ADR-0002), e.g. ``randn_like(mu)`` is noise, not a function of μ.
+SHAPE_ONLY: frozenset[str] = frozenset(
+    _names("""zeros_like ones_like full_like empty_like randn_like rand_like randint_like
+              new_zeros new_ones new_full new_empty new_tensor""")
+)
+
+#: Fused operations that hide sub-module structure (ADR-0003 consequence).
+FUSED: frozenset[str] = frozenset(
+    _names("_transformer_encoder_layer_fwd _native_multi_head_attention")
+)
+
+_REFLECTED = frozenset(
+    _names("add sub mul truediv floordiv mod pow matmul and or xor lshift rshift")
+)
+_RENAMES = {
+    "truediv": "div",
+    "floordiv": "floor_divide",
+    "mod": "remainder",
+    "and": "logical_and",
+    "or": "logical_or",
+    "xor": "logical_xor",
+}
+
+_GROUP_KEYWORDS: tuple[tuple[tuple[str, ...], C], ...] = (
+    (("lstm", "gru", "rnn", "recurrent"), C.RECURRENT),
+    (("attention", "attn"), C.ATTENTION),
+    (("embed",), C.EMBEDDING),
+    (("conv",), C.CONV),
+    (("norm",), C.NORM),
+)
+
+# --------------------------------------------------------------------------- lookups
+
+
+def normalize_op_name(name: str) -> str:
+    """``__add__``/``__radd__`` → ``add``; ``__iadd__`` → ``add_``; ``__getitem__`` →
+    ``getitem``; ``__truediv__`` → ``div``."""
+    if name.startswith("__") and name.endswith("__"):
+        core = name.strip("_")
+        if core.startswith("r") and core[1:] in _REFLECTED:
+            core = core[1:]
+        elif core.startswith("i") and core[1:] in _REFLECTED:
+            return _RENAMES.get(core[1:], core[1:]) + "_"
+        return _RENAMES.get(core, core)
+    return name
+
+
+def base_op(name: str) -> str:
+    """In-place variants share the category of their base op: ``add_`` → ``add``."""
+    return name[:-1] if name.endswith("_") and not name.startswith("_") else name
+
+
+def module_category(mro_names: Iterable[str]) -> C | None:
+    for name in mro_names:
+        if name in MODULE_CATEGORY:
+            return MODULE_CATEGORY[name]
+    return None
+
+
+def module_param_names(mro_names: Iterable[str]) -> tuple[str, ...]:
+    for name in mro_names:
+        if name in MODULE_PARAMS:
+            return MODULE_PARAMS[name]
+    return ()
+
+
+def function_category(name: str, tensor_inputs: int) -> tuple[C, bool]:
+    """Category of a functional op and whether it is known to the catalogue."""
+    op = base_op(name)
+    if op in ARITHMETIC:
+        return (C.MERGE if tensor_inputs >= 2 else C.OTHER), True
+    if op in FUNCTION_CATEGORY:
+        return FUNCTION_CATEGORY[op], True
+    return C.OTHER, op in KNOWN_OTHER
+
+
+def guess_group_category(type_name: str) -> C | None:
+    """Category hint for a user-defined module from its class name (``ConvLSTMCell`` →
+    recurrent). Only used when the module is drawn collapsed."""
+    lower = type_name.lower()
+    for keywords, category in _GROUP_KEYWORDS:
+        if any(k in lower for k in keywords):
+            return category
+    return None