embedl-deploy 0.1.1__tar.gz → 0.2.0__tar.gz

{embedl_deploy-0.1.1/src/embedl_deploy.egg-info → embedl_deploy-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedl-deploy
-Version: 0.1.1
+Version: 0.2.0
 Summary: Python package to make AI models deployment-ready for any hardware.
 Author-email: Embedl AB <support@embedl.com>
 Project-URL: Homepage, https://www.embedl.com/

embedl_deploy-0.2.0/src/embedl_deploy/_internal/core/backend.py

@@ -0,0 +1,125 @@
+# Copyright (C) 2026 Embedl AB
+
+"""Backend discovery and selection."""
+
+import importlib
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+from embedl_deploy._internal.core.pattern import Pattern
+
+_INTERNAL_DIR = Path(__file__).resolve().parent.parent
+
+
+@dataclass(frozen=True)
+class Backend:
+    """A collection of patterns for a specific hardware target."""
+
+    #: Structural rewrite patterns, applied iteratively.
+    conversion_patterns: Sequence[Pattern]
+    #: Fusion patterns, applied in a single pass after conversions.
+    fusion_patterns: Sequence[Pattern]
+    #: SmoothQuant preparation patterns.
+    smooth_patterns: Sequence[Pattern]
+    #: Q/DQ stub insertion patterns for quantisation.
+    quantized_patterns: Sequence[Pattern]
+
+
+class _BackendState:
+    """Module-level mutable state for backend discovery and selection."""
+
+    #: The currently selected backend.
+    backend: Backend | None = None
+    #: Cached discovery result.
+    backends: dict[str, Backend] | None = None
+
+    @classmethod
+    def reset(cls) -> None:
+        """Clear cached discovery results and the active backend."""
+        cls.backend = None
+        cls.backends = None
+
+
+def _discover_backends() -> dict[str, Backend]:
+    """Scan ``_internal/`` for importable backend packages.
+
+    Each subdirectory (except ``core``) is tried as
+    ``embedl_deploy._internal.<name>.backend``.  Directories whose
+    module cannot be found are skipped; import errors from
+    transitive dependencies are propagated.  Results are cached
+    after the first call.
+
+    :returns:
+        Mapping of backend name to ``Backend`` instance.
+    """
+    backends = _BackendState.backends
+    if backends is None:
+        backends = {}
+        for entry in sorted(_INTERNAL_DIR.iterdir()):
+            if (
+                not entry.is_dir()
+                or entry.name.startswith("_")
+                or entry.name == "core"
+            ):
+                continue
+            module_path = f"embedl_deploy._internal.{entry.name}.backend"
+            try:
+                mod = importlib.import_module(module_path)
+            except ModuleNotFoundError as e:
+                if e.name == module_path:
+                    continue
+                raise
+            backend = getattr(mod, "BACKEND", None)
+            if isinstance(backend, Backend):
+                backends[entry.name] = backend
+        _BackendState.backends = backends
+    return backends
+
+
+def get_backend() -> Backend:
+    """Return the active backend, discovering it if necessary.
+
+    If no backend has been set via :func:`set_backend`, the installed
+    backends are discovered automatically.  When exactly one is found
+    it becomes the active backend.
+
+    :returns:
+        The active :class:`Backend`.
+    :raises RuntimeError:
+        If no backends are installed, or if multiple backends are
+        installed and none has been explicitly selected.
+    """
+    backend = _BackendState.backend
+    if backend is None:
+        backends = _discover_backends()
+        if len(backends) == 0:
+            raise RuntimeError(
+                "No backends found — install at least one backend"
+            )
+        if len(backends) > 1:
+            names = ", ".join(sorted(backends))
+            raise RuntimeError(
+                f"Multiple backends found ({names}). "
+                "Call set_backend() to select one."
+            )
+        backend = next(iter(backends.values()))
+        _BackendState.backend = backend
+    return backend
+
+
+def set_backend(name: str) -> None:
+    """Select the active backend by name.
+
+    :param name:
+        The name of a discovered backend (e.g. ``"tensorrt"``).
+    :raises ValueError:
+        If `name` does not match any installed backend.
+    """
+    backends = _discover_backends()
+    if name not in backends:
+        available = ", ".join(sorted(backends)) or "(none)"
+        raise ValueError(
+            f"Backend {name!r} not found. " f"Available backends: {available}"
+        )
+    _BackendState.backend = backends[name]

{embedl_deploy-0.1.1 → embedl_deploy-0.2.0}/src/embedl_deploy/_internal/core/match.py

@@ -9,9 +9,9 @@ occurrences of operation chains in ``torch.fx`` graphs.
 """
 
 import itertools
-from typing import TypeAlias, get_args
+import types
 
-from torch import fx, nn
+from torch import fx
 
 from embedl_deploy._internal.core.pattern import (
     Fork,
@@ -19,6 +19,7 @@ from embedl_deploy._internal.core.pattern import (
     NodeCheck,
     Pattern,
     PatternMatch,
+    SharedNodeCheck,
     Tree,
     TreeMatch,
     Trunk,
@@ -26,28 +27,40 @@ from embedl_deploy._internal.core.pattern import (
     get_module,
 )
 
-ActivationModules: TypeAlias = (
-    nn.ReLU
-    | nn.ReLU6
-    | nn.GELU
-    | nn.SiLU
-    | nn.Mish
-    | nn.Hardswish
-    | nn.Hardsigmoid
-    | nn.LeakyReLU
-    | nn.PReLU
-    | nn.ELU
-    | nn.Sigmoid
-    | nn.Tanh
-)
 
-#: Module types recognized as activation functions by the matchers.
-ACTIVATION_MODULES: tuple[type[nn.Module], ...] = get_args(ActivationModules)
+class _SharedNodeCheckSession:
+    """Checkpoint/rollback scope for ``SharedNodeCheck`` cache entries.
+
+    Between permutation attempts inside a fork, call :meth:`rollback`
+    explicitly.  On normal exit (success) cache entries survive so that
+    enclosing sessions can still enforce cross-fork shared-node
+    constraints.  On failure exit (no permutation matched and the block
+    falls through) ``__exit__`` rolls back automatically.
+    """
+
+    def __init__(self) -> None:
+        self._checkpoint = SharedNodeCheck.checkpoint()
+        self._succeeded = False
+
+    def __enter__(self) -> "_SharedNodeCheckSession":
+        return self
+
+    def success(self) -> None:
+        """Mark this session as successful — skip rollback on exit."""
+        self._succeeded = True
+
+    def __exit__(self, *args: object) -> None:
+        if not self._succeeded:
+            self.rollback()
+
+    def rollback(self) -> None:
+        """Reset every cache event logged since this session started."""
+        SharedNodeCheck.rollback_to(self._checkpoint)
 
 
 def _node_matches(node: fx.Node, checks: ModType | NodeCheck) -> bool:
     """Return whether `node` satisfies `checks`."""
-    if isinstance(checks, (type, tuple)):
+    if isinstance(checks, (type, types.UnionType)):
         return isinstance(get_module(node), checks)
     return checks(node)
 
@@ -133,34 +146,43 @@ def _match_fork_at(
         return None
 
     fork_node = trunk.pre_trunk_nodes[0]
-    if (
+    if getattr(fork.operator, "is_node_check", False):
+        if not fork.operator(fork_node):
+            return None
+    elif (
         fork_node.op != "call_function"
         or fork_node.target is not fork.operator
     ):
         return None
 
     args = [a for a in fork_node.args if isinstance(a, fx.Node)]
-    if len(args) != len(fork.inputs):
-        return None
-
-    for perm in itertools.permutations(range(len(fork.inputs))):
-        tree_matches: dict[int, TreeMatch] = {}
-        for arg_idx, branch_idx in enumerate(perm):
-            matched = _match_tree_at(
-                args[arg_idx],
-                fork.inputs[branch_idx],
-            )
-            if matched is None:
-                break
-            tree_matches[branch_idx] = matched
-
-        if len(tree_matches) == len(fork.inputs):
-            nested = [tree_matches[i] for i in range(len(fork.inputs))]
-            return TreeMatch(
-                pre_trunk_nodes=[fork_node],
-                trunk_nodes=trunk.trunk_nodes,
-                nested=nested,
-            )
+    perms = (
+        itertools.permutations(range(len(fork.inputs)))
+        if fork.perms_override is None
+        else fork.perms_override
+    )
+    with _SharedNodeCheckSession() as session:
+        for perm in perms:
+            session.rollback()
+            if len(perm) != len(args):
+                continue
+            fork_matched = True
+            tree_matches = [TreeMatch() for _ in fork.inputs]
+            for arg_idx, input_idx in enumerate(perm):
+                tree_match = _match_tree_at(
+                    args[arg_idx], fork.inputs[input_idx]
+                )
+                if tree_match is None:
+                    fork_matched = False
+                    break
+                tree_matches[input_idx] = tree_match
+            if fork_matched:
+                session.success()
+                return TreeMatch(
+                    pre_trunk_nodes=[fork_node],
+                    trunk_nodes=trunk.trunk_nodes,
+                    nested=tree_matches,
+                )
     return None
 
 
@@ -218,15 +240,17 @@ def match_tree(
         raise ValueError("``pattern`` has no tree to match.")
 
     matches: list[PatternMatch] = []
-    for node in reversed(graph_module.graph.nodes):
-        matched = _match_tree_at(node, tree)
-        if matched is not None:
-            matches.append(
-                PatternMatch(
-                    pattern=pattern,
-                    graph_module=graph_module,
-                    tree_match=matched,
+    with _SharedNodeCheckSession() as session:
+        for node in reversed(list(graph_module.graph.nodes)):
+            session.rollback()
+            matched = _match_tree_at(node, tree)
+            if matched is not None:
+                matches.append(
+                    PatternMatch(
+                        pattern=pattern,
+                        graph_module=graph_module,
+                        tree_match=matched,
+                    )
                 )
-            )
 
     return matches

{embedl_deploy-0.1.1 → embedl_deploy-0.2.0}/src/embedl_deploy/_internal/core/modules.py

@@ -6,10 +6,26 @@
 # torch lacks type stubs, so nn.Module resolves to Any.
 
 from abc import ABC
+from typing import TypeAlias
 
 from torch import fx, nn
 
-from embedl_deploy._internal.core.quantize.modules import QuantStub
+from embedl_deploy._internal.core.quantize.stubs import QuantStub
+
+ActivationLike: TypeAlias = (
+    nn.ReLU
+    | nn.ReLU6
+    | nn.GELU
+    | nn.SiLU
+    | nn.Mish
+    | nn.Hardswish
+    | nn.Hardsigmoid
+    | nn.LeakyReLU
+    | nn.PReLU
+    | nn.ELU
+    | nn.Sigmoid
+    | nn.Tanh
+)
 
 
 class ConvertedModule(nn.Module, ABC):
@@ -29,13 +45,13 @@ class FusedModule(nn.Module, ABC):
 
     Backend packages (e.g. ``tensorrt``) subclass this for their concrete
     fused modules (``FusedConvBNAct``, etc.).  The generic Q/DQ insertion
-    pass in :mod:`~embedl_deploy._internal.core.quantize.qdq` uses
+    pass in :mod:`~embedl_deploy._internal.core.quantize.prepare` uses
     ``isinstance(mod, FusedModule)`` to identify fused nodes without
     knowing backend-specific types.
     """
 
     #: Positional argument indices that should receive a
-    #: :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub`.
+    #: :class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub`.
     #: The Q/DQ insertion pass uses this to decide which inputs of the
     #: fused node to quantize.  Every subclass must set this explicitly.
     inputs_to_quantize: set[int]
@@ -43,7 +59,7 @@ class FusedModule(nn.Module, ABC):
     def __init__(self) -> None:
         super().__init__()
         #: Maps each index in :attr:`inputs_to_quantize` to a
-        #: :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub`.
+        #: :class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub`.
         self.input_quant_stubs: dict[int, QuantStub] = {
             idx: QuantStub({self}) for idx in self.inputs_to_quantize
         }

{embedl_deploy-0.1.1 → embedl_deploy-0.2.0}/src/embedl_deploy/_internal/core/pattern.py

@@ -9,9 +9,9 @@ methods — :meth:`~embedl_deploy._internal.core.pattern.Pattern.match` and
 what to look for and how to rewrite the graph.
 """
 
-import enum
+import types
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any, Literal, TypedDict, TypeVar, overload
 
@@ -41,7 +41,7 @@ def resolve_module(
     :returns:
         The ``nn.Module`` instance, narrowed to `cls`.
     :raises LookupError:
-        If *node* is not a ``call_module`` node, the target module
+        If `node` is not a ``call_module`` node, the target module
         cannot be found, or the module is not an instance of `cls`.
     """
     if not (
@@ -72,8 +72,8 @@ def get_module(node: fx.Node) -> nn.Module | None:
         return None
 
 
-#: A single module type or a tuple of module types for ``isinstance`` checks.
-ModType = type[nn.Module] | tuple[type[nn.Module], ...]
+#: A single module type or a union of module types for ``isinstance`` checks.
+ModType = type[nn.Module] | types.UnionType
 
 #: A callable predicate that inspects a node and returns whether it matches.
 NodeCheck = Callable[[fx.Node], bool]
@@ -83,11 +83,12 @@ NodeCheck = Callable[[fx.Node], bool]
 class Wildcard:
     """Zero or more consecutive nodes to absorb during matching.
 
-    When a :class:`Wildcard` appears in a
+    When a :class:`~embedl_deploy._internal.core.pattern.Wildcard` appears in a
     :data:`~embedl_deploy._internal.core.pattern.Trunk`, the matcher walks
     backward through nodes that satisfy `checks`.  The absorbed nodes are
-    recorded in `nodes` and the :class:`Wildcard` itself is stored as an
-    entry in
+    recorded in `nodes` and the
+    :class:`~embedl_deploy._internal.core.pattern.Wildcard` itself is stored as
+    an entry in
     :attr:`~embedl_deploy._internal.core.pattern.TreeMatch.trunk_nodes`.
     """
 
@@ -101,7 +102,7 @@ class Wildcard:
 
     @property
     def name(self) -> list[str]:
-        """Node names of absorbed nodes, for use in serialisation."""
+        """Node names of absorbed nodes, for use in serialization."""
         return [n.name for n in self.nodes]
 
 
@@ -110,7 +111,68 @@ class Wildcard:
 #: (for ``isinstance`` checks), a
 #: :data:`~embedl_deploy._internal.core.pattern.NodeCheck` predicate,
 #: or a :class:`~embedl_deploy._internal.core.pattern.Wildcard`.
-Trunk = Sequence[ModType | NodeCheck | Wildcard]
+Trunk = tuple[ModType | NodeCheck | Wildcard, ...]
+
+
+def node_check(
+    fn: Callable[[fx.Node], bool],
+) -> Callable[[fx.Node], bool]:
+    """Mark a function as a ``Fork.operator`` predicate.
+
+    When the matcher encounters a
+    :class:`~embedl_deploy._internal.core.pattern.Fork` whose ``operator``
+    carries the ``is_node_check`` attribute, it calls the operator with the
+    candidate node and expects a ``bool``. Without the marker the operator is
+    identity-compared against ``node.target``.
+
+    :param fn:
+        A :data:`~embedl_deploy._internal.core.pattern.NodeCheck` callable.
+    :returns:
+        The same callable, with ``is_node_check = True`` attached.
+    """
+    setattr(fn, "is_node_check", True)
+    return fn
+
+
+class SharedNodeCheck:
+    """Constrain multiple fork branches to the same physical node.
+
+    On the first call, delegates to an inner ``NodeCheck`` and caches the
+    matched node.  Subsequent calls succeed only when the candidate is the
+    exact same ``fx.Node`` instance.  A class-level log records every
+    cache event so the matcher can checkpoint and rollback.
+    """
+
+    _log: list["SharedNodeCheck"] = []
+
+    def __init__(self, inner: NodeCheck) -> None:
+        self._inner = inner
+        self._cached: fx.Node | None = None
+
+    def __call__(self, node: fx.Node) -> bool:
+        if self._cached is None:
+            if not self._inner(node):
+                return False
+            self._cached = node
+            SharedNodeCheck._log.append(self)
+            return True
+        return node is self._cached
+
+    def reset(self) -> None:
+        """Clear the cached node."""
+        self._cached = None
+
+    @classmethod
+    def checkpoint(cls) -> int:
+        """Return a rollback point for use with ``rollback_to``."""
+        return len(cls._log)
+
+    @classmethod
+    def rollback_to(cls, checkpoint: int) -> None:
+        """Reset and discard every log entry from `checkpoint` onward."""
+        for entry in cls._log[checkpoint:]:
+            entry.reset()
+        del cls._log[checkpoint:]
 
 
 @dataclass(frozen=True)
@@ -125,12 +187,71 @@ class Fork:
     #: is itself a :data:`~embedl_deploy._internal.core.pattern.Tree`
     #: (a :data:`~embedl_deploy._internal.core.pattern.Trunk` or
     #: another :class:`~embedl_deploy._internal.core.pattern.Fork`).
-    #: An empty sequence means the branch accepts any single node.
-    inputs: Sequence["Tree"]
-    #: The function that the merge node must call (e.g. ``operator.add``).
+    #: An empty tuple means the branch accepts any single node.
+    inputs: tuple["Tree", ...]
+    #: The function that the fork node must invoke.  There are two
+    #: modes, selected by the presence of the ``is_node_check``
+    #: attribute (set by the :func:`node_check` decorator):
+    #:
+    #: **Identity compare** (default) — the fork node's ``target``
+    #: is checked with ``is`` against this value.  Use for concrete
+    #: functions like ``operator.add`` or
+    #: ``torch.ops.aten.conv2d.default``.
+    #:
+    #: **Predicate** (``@node_check``) — called with the fork node;
+    #: must return ``True`` for the match to succeed.  Use when a
+    #: single function target is not sufficient, e.g. when the match
+    #: also depends on argument shapes or metadata.
     operator: Callable[..., Any]
     #: The straight chain of modules after the merge node.
     output: Trunk
+    #: Optional override of the permitted mappings from the arguments of the
+    #: ``operator`` to the branches in ``inputs``. Only ``fx.Node`` arguments
+    #: are considered (``None`` and scalar arguments are ignored). Each
+    #: mapping is a tuple whose length equals the expected number of such
+    #: arguments. At index *i*, a value of *k* in a mapping means that the
+    #: *i*-th argument of the ``operator`` must match the *k*-th branch in
+    #: ``inputs``. A mapping may not be longer than ``inputs``, but it may be
+    #: shorter if some of the branches are optional.
+    #:
+    #: .. code-block:: python
+    #:
+    #:     fork = Fork(
+    #:         inputs=(
+    #:             (),                     # input tensor
+    #:             (_is_attr("weight"),),
+    #:             (_is_attr("bias"),),    # optional
+    #:         ),
+    #:         operator=torch.ops.aten.conv2d.default,
+    #:         output=(),
+    #:         perms_override=((0, 1, 2), (0, 1)),
+    #:     )
+    #:
+    #: The two entries in ``perms_override`` cover the two shapes of the
+    #: ``aten.conv2d.default`` call: ``(0, 1, 2)`` matches when bias is
+    #: present, ``(0, 1)`` matches when it is absent.
+    #:
+    #: By default the ``perms_override`` is ``None``. In that case, the
+    #: arguments of the ``operator`` are checked against all full-length
+    #: permutations of the branches in ``inputs`` until one matches.
+    perms_override: tuple[tuple[int, ...], ...] | None = None
+
+    def __post_init__(self) -> None:
+        if self.perms_override is None:
+            return
+        valid_set = set(range(len(self.inputs)))
+        for perm in self.perms_override:
+            perm_set = set(perm)
+            if not perm_set <= valid_set:
+                raise ValueError(
+                    f"perms_override entry {perm!r} contains "
+                    f"out-of-range indices"
+                )
+            if len(perm_set) < len(perm):
+                raise ValueError(
+                    f"perms_override entry {perm!r} contains "
+                    f"duplicate indices"
+                )
 
 
 #: A pattern topology: either a straight
@@ -139,24 +260,6 @@ class Fork:
 Tree = Trunk | Fork
 
 
-class QDQPoint(enum.Enum):
-    """Named positions where quantization stubs are inserted.
-
-    Each :class:`~embedl_deploy._internal.core.pattern.Pattern` subclass
-    declares which Q/DQ points it needs via its
-    :attr:`~embedl_deploy._internal.core.pattern.Pattern.qdq_points` class
-    attribute.  The Q/DQ insertion pass uses these declarations to place
-    :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub` nodes at
-    the correct positions in the graph.
-    """
-
-    INPUT = "input"
-    OUTPUT = "output"
-    RESIDUAL_INPUT = "residual_input"
-    KEY_INPUT = "key_input"
-    VALUE_INPUT = "value_input"
-
-
 class Pattern(ABC):
     """A graph transformation rule: find a sub-graph and replace it.
 
@@ -179,12 +282,10 @@ class Pattern(ABC):
     """If ``True``, this pattern is a structural conversion that must
     be applied before fusion matching."""
 
-    qdq_points: frozenset[QDQPoint] = frozenset()
-    """Q/DQ insertion points declared by this pattern.
-
-    Subclasses override this to declare where
-    :class:`~embedl_deploy._internal.core.quantize.modules.QuantStub`
-    nodes should be placed.  The default (empty) means no Q/DQ stubs."""
+    symbolic_trace_only: bool = False
+    """If ``True``, this pattern removes nodes that are artifacts of
+    ``symbolic_trace``. This pattern has no effect on graphs exported with
+    ``torch.export`` because the nodes never appear in those graphs."""
 
     @abstractmethod
     def match(self, graph_module: fx.GraphModule) -> list["PatternMatch"]:
@@ -215,10 +316,10 @@ class Pattern(ABC):
         """
 
 
-class SerialisedTreeMatch(TypedDict):
-    """JSON-serialisable representation of a :class:`TreeMatch`."""
+class SerializedTreeMatch(TypedDict):
+    """JSON-serializable representation of a ``TreeMatch``."""
 
-    nested: list["SerialisedTreeMatch"]
+    nested: list["SerializedTreeMatch"]
     pre_trunk_nodes: list[str]
     trunk_nodes: list[str | list[str]]
 
@@ -247,6 +348,7 @@ class TreeMatch:
     trunk_nodes: list[fx.Node | Wildcard] = field(default_factory=list)
     #: Per-branch sub-matches, indexed the same as
     #: :attr:`~embedl_deploy._internal.core.pattern.Fork.inputs`.
+    #: Empty entries represent optional branches that were absent.
     nested: list["TreeMatch"] = field(default_factory=list)
 
     def get_input_nodes(self) -> tuple[fx.Node, ...]:
@@ -311,9 +413,10 @@ class TreeMatch:
 
         With a single `index`, returns ``trunk_nodes[index]``.  For a
         :class:`~embedl_deploy._internal.core.pattern.Wildcard` position the
-        returned value is the :class:`Wildcard` itself (access its ``nodes``
-        attribute for the absorbed nodes).  With additional indices the lookup
-        descends into
+        returned value is the
+        :class:`~embedl_deploy._internal.core.pattern.Wildcard` itself
+        (access its ``nodes`` attribute for the absorbed nodes).  With
+        additional indices the lookup descends into
         :attr:`~embedl_deploy._internal.core.pattern.TreeMatch.nested`
         sub-matches first.
 
@@ -325,10 +428,12 @@ class TreeMatch:
         :param args:
             Additional indices forwarded recursively into nested sub-matches.
         :param is_wildcard:
-            If ``True``, assert the entry is a :class:`Wildcard` and return
+            If ``True``, assert the entry is a
+            :class:`~embedl_deploy._internal.core.pattern.Wildcard` and return
             it.  If ``False`` (default), assert it is an ``fx.Node``.
         :returns:
-            The matched ``fx.Node`` or :class:`Wildcard`.
+            The matched ``fx.Node`` or
+            :class:`~embedl_deploy._internal.core.pattern.Wildcard`.
         """
         if args:
             nested = self.nested[index]
@@ -339,10 +444,10 @@ class TreeMatch:
         assert isinstance(entry, Wildcard if is_wildcard else fx.Node)
         return entry
 
-    def serialise(self) -> SerialisedTreeMatch:
-        """Serialise this match to a JSON-serialisable dict."""
+    def serialize(self) -> SerializedTreeMatch:
+        """Serialize this match to a JSON-serializable dict."""
         return {
-            "nested": [nm.serialise() for nm in self.nested],
+            "nested": [nm.serialize() for nm in self.nested],
             "pre_trunk_nodes": [n.name for n in self.pre_trunk_nodes],
             "trunk_nodes": [n.name for n in self.trunk_nodes],
         }