# Source code for tensorplay.export._trace

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/export/_trace.html

```
"""Capture routines for building structured exported programs."""

from __future__ import annotations

import inspect
from collections.abc import Mapping
from typing import Any, Callable

from ..graph import GraphCaptureError, GraphModule, Node, Proxy, Tracer
from .dynamic_shapes import AdditionalInputs, ConstraintsExceededError, Dim, ShapesCollection, _DimHint
from .exported_program import EqualityConstraint, ExportedProgram, ModuleCallEntry, ModuleCallSignature
from .graph_signature import (
    ConstantArgument,
    ExportGraphSignature,
    InputKind,
    InputSpec,
    OutputKind,
    OutputSpec,
    TensorArgument,
)

__all__ = ["ExportTracer", "draft_export", "export", "export_for_training"]

_STATE_PREFIXES = {"parameter": "p", "buffer": "b", "constant": "c"}

def _qualified(module_name: str, attribute: str) -> str:
    return f"{module_name}.{attribute}" if module_name else attribute

def _collect_attributes(root: Any) -> dict[str, tuple[str, bool]]:
    """Map qualified attribute paths to ``(kind, persistent)``.

    Kinds cover parameters, persistent and non-persistent buffers, and plain
    tensor attributes (recorded as constants).  Constant entries carry the
    value so the tracer can lift them alongside parameters and buffers.
    """

    attributes: dict[str, tuple[str, bool]] = {}
    named_modules = getattr(root, "named_modules", None)
    if not callable(named_modules):
        return attributes
    import tensorplay as tp

    for module_name, module in named_modules(remove_duplicate=True):
        non_persistent = getattr(module, "_non_persistent_buffers_set", set())
        for name, value in getattr(module, "_parameters", {}).items():
            if value is not None:
                attributes[_qualified(module_name, name)] = ("parameter", True)
        for name, value in getattr(module, "_buffers", {}).items():
            if value is not None:
                attributes[_qualified(module_name, name)] = (
                    "buffer",
                    name not in non_persistent,
                )
        parameters = getattr(module, "_parameters", {})
        buffers = getattr(module, "_buffers", {})
        children = getattr(module, "_modules", {})
        for name, value in vars(module).items():
            if name.startswith("_") or name in parameters or name in buffers:
                continue
            if name in children or callable(value):
                continue
            if isinstance(value, tp.nn.Parameter) or not isinstance(value, tp.Tensor):
                continue
            attributes[_qualified(module_name, name)] = ("constant", True)
    return attributes

def _resolve_attribute(root: Any, target: str) -> Any:
    value = root
    for atom in target.split("."):
        value = getattr(value, atom)
    return value

class ExportTracer(Tracer):
    """Tracer that lifts module state into graph inputs.

    Parameters, buffers, and constant tensors become placeholders ahead of the
    user inputs, so the captured graph is functional: it reads no attributes
    and its only inputs are the flat value list described by the graph
    signature.

    Child module forwards are additionally wrapped so every recorded node
    carries the qualified path of the module that produced it
    (``nn_module_stack``), and module call boundaries (argument and result
    nodes) are recorded for later hierarchy reconstruction.
    """

    def __init__(self, concrete_args: dict[str, Any] | None = None) -> None:
        super().__init__(concrete_args)
        # qualified attribute path -> (placeholder node, kind, persistent)
        self.state_targets: dict[str, tuple[Node, str, bool]] = {}
        self._constant_patches: list[tuple[Any, str, Any]] = []
        self._missing_sentinel = object()
        self._module_stack: tuple[str, ...] = ()
        self._forward_patches: list[tuple[Any, Any]] = []
        self.module_calls: list[dict[str, Any]] = []
        self._call_keys: set[tuple[Any, ...]] = set()

    def _register_state(self, root: Any) -> None:
        for target, (kind, persistent) in _collect_attributes(root).items():
            value = _resolve_attribute(root, target)
            mangled = f"{_STATE_PREFIXES[kind]}_{target.replace('.', '_')}"
            node = self.graph.create_node("placeholder", mangled, (), {}, name=mangled)
            node.meta["state_target"] = target
            node.meta["state_kind"] = kind
            node.meta["state_persistent"] = persistent
            self.state_targets[target] = (node, kind, persistent)

    def _patch_constants(self, root: Any) -> None:
        for target, (node, kind, _persistent) in self.state_targets.items():
            if kind != "constant":
                continue
            parent_name, _, leaf = target.rpartition(".")
            parent = _resolve_attribute(root, parent_name) if parent_name else root
            previous = getattr(parent, leaf, self._missing_sentinel)
            self._constant_patches.append((parent, leaf, previous))
            setattr(parent, leaf, Proxy(node, self))

    def _restore_constants(self) -> None:
        for parent, leaf, previous in reversed(self._constant_patches):
            if previous is self._missing_sentinel:
                delattr(parent, leaf)
            else:
                setattr(parent, leaf, previous)
        self._constant_patches.clear()

    def _wrap_child_forwards(self, root: Any) -> None:
        """Route child module calls through stack-tracking wrappers."""

        from ..graph._pytree import tree_flatten

        named_modules = getattr(root, "named_modules", None)
        if not callable(named_modules):
            return
        for module_name, module in named_modules(remove_duplicate=True):
            if not module_name:
                continue  # the root call is described by the program itself
            original = getattr(module, "forward", None)
            if not callable(original):
                continue

            def wrapper(
                *args: Any,
                _original: Any = original,
                _qualname: str = module_name,
                **kwargs: Any,
            ) -> Any:
                arg_nodes: list[Any] = []
                for value in args:
                    arg_nodes.append(
                        value.node.name if isinstance(value, Proxy) else value
                    )
                kwargs_nodes = {
                    key: value.node.name if isinstance(value, Proxy) else value
                    for key, value in kwargs.items()
                }
                self._module_stack = (*self._module_stack, _qualname)
                try:
                    result = _original(*args, **kwargs)
                finally:
                    self._module_stack = self._module_stack[:-1]
                result_nodes: list[str] = []

                def visit(item: Any) -> None:
                    if isinstance(item, Proxy):
                        result_nodes.append(item.node.name)
                    elif isinstance(item, (tuple, list)):
                        for entry in item:
                            visit(entry)
                    elif isinstance(item, dict):
                        for entry in item.values():
                            visit(entry)

                visit(result)
                _in_spec = tree_flatten(args)[1]
                _out_spec = tree_flatten(result)[1]
                key = (_qualname, tuple(map(repr, arg_nodes)), tuple(result_nodes))
                if key not in self._call_keys:
                    self._call_keys.add(key)
                    self.module_calls.append(
                        {
                            "fqn": _qualname,
                            "args": arg_nodes,
                            "kwargs": kwargs_nodes,
                            "result": result_nodes,
                            "in_spec": _in_spec,
                            "out_spec": _out_spec,
                        }
                    )
                return result

            try:
                module.forward = wrapper  # type: ignore[method-assign]
            except Exception:
                continue
            self._forward_patches.append((module, original))

    def _restore_child_forwards(self) -> None:
        for module, _original in self._forward_patches:
            try:
                del module.forward
            except AttributeError:
                pass
        self._forward_patches.clear()

    def trace(self, root: Any, sample_inputs: dict[str, Any] | None = None) -> GraphModule:
        self.root = root
        if callable(getattr(root, "named_modules", None)):
            self._register_state(root)
            self._patch_constants(root)
            self._wrap_child_forwards(root)
            try:
                return super().trace(root, sample_inputs)
            finally:
                self._restore_child_forwards()
                self._restore_constants()
        return super().trace(root, sample_inputs)

    def create_proxy(
        self,
        kind: str,
        target: Any,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
    ) -> Proxy:
        if kind == "get_attr":
            entry = self.state_targets.get(target)
            if entry is not None:
                return Proxy(entry[0], self)
        proxy = super().create_proxy(kind, target, args, kwargs)
        if self._module_stack:
            proxy.node.meta["nn_module_stack"] = self._module_stack
        return proxy

def _validate_graph(graph_module: GraphModule, attributes: Mapping[str, Any]) -> None:
    for node in graph_module.graph.nodes:
        op = node.op
        if op == "get_attr":
            # A tensor the function closed over is not an attribute of the
            # model; the capture holds it itself, under a name of its own.
            if (
                node.target not in attributes
                and node.target not in graph_module._graph_attrs
            ):
                raise GraphCaptureError(
                    f"get_attr target {node.target!r} is not present on the captured model"
                )
        elif op == "call_function":
            if not callable(node.target):
                raise GraphCaptureError(f"call_function target {node.target!r} is not callable")
        elif op == "call_method":
            if not node.args or not isinstance(node.target, str):
                raise GraphCaptureError(f"malformed call_method node: {node}")
        elif op not in {"placeholder", "output", "call_module"}:
            raise GraphCaptureError(f"unsupported graph node kind: {op!r}")
    if not graph_module.graph.outputs:
        raise GraphCaptureError("captured graph has no output")

def _normalize_dynamic_shapes(
    spec: Any,
    parameter_names: list[str],
    model: Any,
    args: tuple[Any, ...],
    kwargs: Mapping[str, Any],
) -> Any:
    if spec is None:
        return {}
    if isinstance(spec, AdditionalInputs):
        return spec.dynamic_shapes(model, args, kwargs)
    if isinstance(spec, ShapesCollection):
        return spec.dynamic_shapes(model, args, kwargs)
    if isinstance(spec, (list, tuple)):
        if len(spec) > len(parameter_names):
            raise ValueError("dynamic_shapes has more entries than graph inputs")
        return {
            name: _validate_dimension_spec(value)
            for name, value in zip(parameter_names, spec)
            if value is not None
        }
    if not isinstance(spec, dict):
        raise TypeError("dynamic_shapes must be a mapping or a sequence")

    normalized: dict[str, Any] = {}
    for name, dims_spec in spec.items():
        if name not in parameter_names:
            raise ValueError(
                f"dynamic_shapes key {name!r} does not match any argument; "
                f"expected one of {parameter_names}"
            )
        if isinstance(dims_spec, dict):
            entry: dict[int, Any] = {}
            for index, value in dims_spec.items():
                if type(index) is not int or index < 0:
                    raise TypeError(f"dimension index must be a non-negative int, got {index!r}")
                entry[index] = _validate_dimension_value(value)
            normalized[name] = entry
        elif isinstance(dims_spec, (tuple, list)):
            normalized[name] = _validate_dimension_spec(dims_spec)
        elif dims_spec is None:
            normalized[name] = None
        else:
            raise TypeError(f"dynamic_shapes[{name!r}] must describe dimensions")
    return normalized

def _validate_dimension_value(value: Any) -> Any:
    if value is None or isinstance(value, (Dim, _DimHint)):
        return value
    if type(value) is int:
        return value
    raise TypeError(
        "dimension spec must be int or Dim (a dim hint or None is also accepted), "
        f"got {type(value)!r}"
    )

def _validate_dimension_spec(value: Any) -> Any:
    if isinstance(value, dict):
        result: dict[int, Any] = {}
        for index, item in value.items():
            if type(index) is not int or index < 0:
                raise TypeError(f"dimension index must be a non-negative int, got {index!r}")
            result[index] = _validate_dimension_value(item)
        return result
    if isinstance(value, (tuple, list)):
        return tuple(_validate_dimension_value(item) for item in value)
    return _validate_dimension_value(value)

def _argument_for_node(node: Any) -> Any:
    if isinstance(node, Node):
        return TensorArgument(node.name)
    if isinstance(node, (str, int, float, bool)) or node is None:
        return ConstantArgument(f"constant_{abs(hash(repr(node))) % 100000}", node)
    return ConstantArgument(f"constant_{abs(hash(type(node).__name__)) % 100000}", node)

def _flatten_leaves(value: Any) -> list[Any]:
    leaves: list[Any] = []

    def visit(item: Any) -> None:
        if isinstance(item, Node):
            leaves.append(item)
        elif isinstance(item, (tuple, list)):
            for entry in item:
                visit(entry)
        elif isinstance(item, dict):
            for entry in item.values():
                visit(entry)
        else:
            leaves.append(item)

    visit(value)
    return leaves

def _mutation_chain_root(node: Node) -> Node:
    """Walk in-place op and element-read chains back to the state holder.

    Two hop rules apply: in-place methods (``add_``) consume the previous
    value of the object they update, and ``getitem`` reads reach into a
    container that itself entered the graph as one placeholder.  Hopping
    through both attributes an element update to the container input.
    """

    import operator

    seen: set[int] = set()
    current = node
    while id(current) not in seen:
        seen.add(id(current))
        if (
            current.op == "call_method"
            and isinstance(current.target, str)
            and current.target.endswith("_")
            and current.args
            and isinstance(current.args[0], Node)
        ):
            current = current.args[0]
            continue
        if (
            current.op == "call_function"
            and current.target is operator.getitem
            and current.args
            and isinstance(current.args[0], Node)
        ):
            current = current.args[0]
            continue
        break
    return current

def _detect_mutations(
    graph_module: GraphModule,
    state_targets: Mapping[str, tuple[Node, str, bool]],
) -> list[tuple[Node, OutputKind, str | None]]:
    """Find in-place updates of lifted state or user inputs.

    Returns ``(holder, kind, target)`` tuples where ``holder`` is the graph
    node carrying the final value of the mutated object.
    """

    state_by_placeholder: dict[str, tuple[str, str, bool]] = {
        node.name: (target, kind, persistent)
        for target, (node, kind, persistent) in state_targets.items()
    }
    holders: dict[str, tuple[OutputKind, str | None]] = {}
    for node in graph_module.graph.placeholders:
        entry = state_by_placeholder.get(node.name)
        if entry is not None:
            _target, kind, _persistent = entry
            holders[node.name] = (
                (
                    OutputKind.PARAMETER_MUTATION
                    if kind == "parameter"
                    else OutputKind.BUFFER_MUTATION
                ),
                _target,
            )
        else:
            holders[node.name] = (OutputKind.USER_INPUT_MUTATION, None)
    mutations: dict[str, tuple[Node, OutputKind, str | None]] = {}
    for node in graph_module.graph.nodes:
        if node.op != "call_method" or not isinstance(node.target, str):
            continue
        if not node.target.endswith("_") or not node.args or not isinstance(node.args[0], Node):
            continue
        root = _mutation_chain_root(node)
        info = holders.get(root.name)
        if info is None:
            continue
        # only a container element (getitem chain) may update a user input;
        # a direct in-place call on a whole input tensor is graph-invisible
        # by construction because capture runs on value copies
        if info[0] is OutputKind.USER_INPUT_MUTATION and root is node.args[0]:
            continue
        target = info[1] if info[1] is not None else root.name
        mutations[node.name] = (node, info[0], target)
    return list(mutations.values())

def _state_entry_for_placeholder(
    state_targets: Mapping[str, tuple[Node, str, bool]],
    placeholder_name: str,
) -> tuple[str, str, bool] | None:
    for target, (node, kind, persistent) in state_targets.items():
        if node.name == placeholder_name:
            return target, kind, persistent
    return None

def _state_map_for_placeholders(
    state_targets: Mapping[str, tuple[Node, str, bool]],
) -> dict[str, tuple[str, str, bool]]:
    """Placeholder-name-keyed view of the lifted-state table.

    Signature construction visits placeholders in graph order; a name-keyed
    map keeps that walk linear instead of scanning the state table per node.
    """

    return {
        node.name: (target, kind, persistent)
        for target, (node, kind, persistent) in state_targets.items()
    }

def _rewrite_container_reads(
    graph_module: GraphModule,
    mutations: list[tuple[Node, OutputKind, str | None]],
) -> None:
    """Point element reads after a mutation at the mutated value.

    A second ``items[0]`` read records its own getitem node; without this
    rewrite it would observe the pre-mutation element and diverge from eager
    execution, where both reads return the same object.
    """

    import operator

    if not mutations:
        return
    graph = graph_module.graph
    order = {node.name: position for position, node in enumerate(graph.nodes)}
    finals: dict[tuple[str, int], Node] = {}
    for node, _kind, _target in mutations:
        current = node
        container: Node | None = None
        index: int | None = None
        while True:
            if (
                current.op == "call_function"
                and current.target is operator.getitem
                and current.args
                and isinstance(current.args[0], Node)
            ):
                if index is None and len(current.args) > 1 and isinstance(current.args[1], int):
                    container = current.args[0]
                    index = current.args[1]
                current = current.args[0]
                continue
            break
        if container is not None and index is not None:
            finals[(container.name, index)] = node
    if not finals:
        return
    for read in list(graph.nodes):
        if read.op != "call_function" or read.target is not operator.getitem:
            continue
        if len(read.args) != 2 or not isinstance(read.args[0], Node):
            continue
        key = (read.args[0].name, read.args[1]) if isinstance(read.args[1], int) else None
        final = finals.get(key) if key is not None else None
        if final is None or final is read:
            continue
        if order.get(final.name, -1) >= order.get(read.name, 1 << 30):
            # the read happens before the mutation; it keeps the old value
            continue
        read.replace_all_uses_with(final)
        if not read.users:
            graph.erase_node(read)

def _output_specs(
    graph_module: GraphModule,
    mutations: list[tuple[Node, OutputKind, str | None]],
) -> list[OutputSpec]:
    output = graph_module.graph.output_node
    leaves = _flatten_leaves(output.args[0])
    specs = [
        OutputSpec(kind, TensorArgument(node.name), target)
        for node, kind, target in mutations
    ]
    specs.extend(
        OutputSpec(OutputKind.USER_OUTPUT, _argument_for_node(value))
        for value in leaves[len(mutations):]
    )
    return specs

def _restructure_output(
    graph_module: GraphModule,
    mutations: list[tuple[Node, OutputKind, str | None]],
) -> None:
    """Rewrite the graph result to ``(mutations..., flattened_user_outputs...)``.

    A mutated value may also be a user output; the flat contract repeats the
    node in both roles, exactly as the signature describes.
    """

    if not mutations:
        return
    graph = graph_module.graph
    output = graph.output_node
    leaves = _flatten_leaves(output.args[0])
    mutation_nodes = [node for node, _kind, _target in mutations]
    graph.output(tuple([*mutation_nodes, *leaves]))

# -- runtime assertions for dynamic-shape contracts ------------------------

def _assert_dim_range(tensor: Any, index: int, min: Any, max: Any, name: str) -> Any:
    size = tuple(tensor.shape)[index]
    if (min is not None and size < min) or (max is not None and size > max):
        raise ConstraintsExceededError(
            f"runtime assertion failed for {name!r}: expected dimension {index} "
            f"in [{min if min is not None else '-inf'}, {max if max is not None else 'inf'}], "
            f"got {size}"
        )
    return size

def _assert_dims_equal(tensor_a: Any, index_a: int, tensor_b: Any, index_b: int, name: str) -> Any:
    size_a = tuple(tensor_a.shape)[index_a]
    size_b = tuple(tensor_b.shape)[index_b]
    if size_a != size_b:
        raise ConstraintsExceededError(
            f"runtime assertion failed for {name!r}: dimensions "
            f"{index_a} and {index_b} must agree, got {size_a} and {size_b}"
        )
    return size_a

def _assert_dim_relation(
    tensor_root: Any,
    index_root: int,
    tensor_derived: Any,
    index_derived: int,
    scale: int,
    offset: int,
    name: str,
) -> Any:
    size_root = tuple(tensor_root.shape)[index_root]
    size_derived = tuple(tensor_derived.shape)[index_derived]
    expected = scale * size_root + offset
    if size_derived != expected:
        raise ConstraintsExceededError(
            f"runtime assertion failed for {name!r}: expected dimension "
            f"{index_derived} == {scale} * dim {index_root} + {offset} "
            f"({expected}), got {size_derived}"
        )
    return size_derived

def _apply_dynamic_shape_constraints(
    graph_module: GraphModule,
    combined_args: Mapping[str, Any],
    normalized: Mapping[str, Any],
) -> tuple[dict[str, dict[str, int | None]], list[EqualityConstraint]]:
    """Validate the spec, insert runtime assertions, and describe shared dims.

    Returns the per-name range bounds and one equality record per named
    dimension that appears at more than one input site (shared dims must hold
    equal sizes across those sites at runtime).
    """

    from .dynamic_shapes import _constraint_program, _process_dynamic_shapes
    from .dim_constraints import DimConstraints, attach_observed_sizes

    constraints = _process_dynamic_shapes(combined_args, normalized)
    if not constraints:
        return {}, []

    # strict checks first: contradictory range declarations for one name are
    # specification errors regardless of the example inputs
    asserts, ranges = _constraint_program(constraints)

    solver = DimConstraints()
    for constraint, observed in attach_observed_sizes(constraints, combined_args, normalized):
        solver.add(constraint, observed)
    if not solver.solve():
        raise ConstraintsExceededError(
            "export-time dimension constraints are inconsistent:\n"
            + solver.pretty_print()
        )

    placeholders = {node.name: node for node in graph_module.graph.placeholders}
    identity_to_name = {
        id(value): name
        for name, value in combined_args.items()
        if not isinstance(value, (dict, list, tuple))
    }

    # one equality record per name spanning several distinct input sites
    site_pairs: dict[str, set[tuple[str, int]]] = {}
    for constraint in constraints:
        if constraint.name is None:
            continue
        input_name = identity_to_name.get(id(constraint.source))
        if input_name is None:
            continue
        site_pairs.setdefault(constraint.name, set()).add((input_name, constraint.dim))
    equality_constraints = [
        EqualityConstraint(tuple(sorted(pairs, key=repr)), name=name)
        for name, pairs in site_pairs.items()
        if len(pairs) > 1
    ]

    def node_of(source: Any) -> Node | None:
        name = identity_to_name.get(id(source))
        return placeholders.get(name) if name is not None else None

    anchors: dict[str, tuple[Node, int]] = {}
    graph = graph_module.graph
    output = graph.output_node
    with graph.inserting_before(output):
        for constraint in asserts:
            if constraint.root is None and constraint.name is not None:
                node = node_of(constraint.source)
                if node is not None:
                    anchors[constraint.name] = (node, constraint.dim)
        for constraint in asserts:
            node = node_of(constraint.source)
            if node is None:
                continue
            if constraint.root is not None:
                anchor = anchors.get(constraint.root)
                if anchor is None:
                    continue
                graph.call_function(
                    _assert_dim_relation,
                    (
                        anchor[0],
                        anchor[1],
                        node,
                        constraint.dim,
                        constraint.scale,
                        constraint.offset,
                        constraint.name or constraint.root,
                    ),
                )
                continue
            if constraint.name is not None:
                anchor = anchors.get(constraint.name)
                if anchor is not None and anchor[0] is not node:
                    graph.call_function(
                        _assert_dims_equal,
                        (
                            anchor[0],
                            anchor[1],
                            node,
                            constraint.dim,
                            constraint.name,
                        ),
                    )
                    continue
            if constraint.min is None and constraint.max is None:
                continue
            graph.call_function(
                _assert_dim_range,
                (
                    node,
                    constraint.dim,
                    constraint.min,
                    constraint.max,
                    constraint.name or f"dim {constraint.dim}",
                ),
            )
    return ranges, equality_constraints

def _graph_signature(
    graph_module: GraphModule,
    state_targets: Mapping[str, tuple[Node, str, bool]],
    mutations: list[tuple[Node, OutputKind, str | None]],
) -> ExportGraphSignature:
    """Signature over the flat input contract: state first, then user inputs."""

    state_by_placeholder = _state_map_for_placeholders(state_targets)
    inputs: list[InputSpec] = []
    for node in graph_module.graph.placeholders:
        entry = state_by_placeholder.get(node.name)
        if entry is None:
            inputs.append(InputSpec(InputKind.USER_INPUT, TensorArgument(node.name), None))
            continue
        target, kind, persistent = entry
        if kind == "parameter":
            inputs.append(
                InputSpec(InputKind.PARAMETER, TensorArgument(node.name), target)
            )
        elif kind == "buffer":
            inputs.append(
                InputSpec(
                    InputKind.BUFFER,
                    TensorArgument(node.name),
                    target,
                    persistent=persistent,
                )
            )
        else:
            inputs.append(
                InputSpec(
                    InputKind.CONSTANT_TENSOR,
                    TensorArgument(node.name),
                    target,
                    persistent=None,
                )
            )
    return ExportGraphSignature(inputs, _output_specs(graph_module, mutations))

def _bind_examples(graph_module: GraphModule, args: tuple[Any, ...], kwargs: Mapping[str, Any]) -> dict[str, Any]:
    signature = graph_module.signature
    if signature is None:
        signature = inspect.signature(getattr(graph_module.root, "forward", graph_module.root))
    bound = signature.bind_partial(*args, **dict(kwargs))
    bound.apply_defaults()
    return {
        node.name: bound.arguments[
            node.target if isinstance(node.target, str) else node.name
        ]
        for node in graph_module.graph.placeholders
        if (
            node.target if isinstance(node.target, str) else node.name
        ) in bound.arguments
    }

def _flat_signature(
    graph_module: GraphModule,
    state_targets: Mapping[str, tuple[Node, str, bool]],
    example_inputs: Mapping[str, Any] | None = None,
) -> inspect.Signature:
    """Signature covering every flat input: state placeholders, then user args.

    User parameters keep their kinds; defaults prefer the export-time binding
    over the source declaration so omitted call arguments replay the capture.
    """

    example_inputs = example_inputs or {}
    state_names = {node.name for node, _kind, _persistent in state_targets.values()}
    user_parameters = (
        dict(graph_module.signature.parameters)
        if graph_module.signature is not None
        else {}
    )
    parameters: list[inspect.Parameter] = []
    for node in graph_module.graph.placeholders:
        if node.name in state_names:
            parameters.append(
                inspect.Parameter(node.name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
            )
            continue
        original = user_parameters.get(node.name)
        default = original.default if original is not None else inspect.Parameter.empty
        if node.name in example_inputs:
            default = example_inputs[node.name]
        if original is not None and default is not original.default:
            original = inspect.Parameter(
                original.name, original.kind, default=default
            )
        if original is not None:
            parameters.append(original)
        elif node.args:
            parameters.append(
                inspect.Parameter(
                    node.name,
                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
                    default=node.args[0],
                )
            )
        else:
            parameters.append(
                inspect.Parameter(node.name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
            )
    return inspect.Signature(parameters)

def _capture(
    model: Callable[..., Any],
    args: tuple[Any, ...],
    kwargs: Mapping[str, Any],
    dynamic_shapes: Any,
) -> ExportedProgram:
    from ..graph._pytree import tree_flatten
    from .dynamic_shapes import _combine_args

    tracer = ExportTracer()
    target = model.forward if callable(getattr(model, "forward", None)) else model
    try:
        bound = inspect.signature(target).bind_partial(*args, **kwargs)
        bound.apply_defaults()
        sample_inputs = dict(bound.arguments)
    except (TypeError, ValueError):
        sample_inputs = None
    from tensorplay.compiler import _exporting_context

    with _exporting_context():
        graph_module = tracer.trace(model, sample_inputs=sample_inputs)
    attributes = _collect_attributes(model)
    _validate_graph(graph_module, attributes)
    placeholders = graph_module.graph.placeholders
    names = [node.name for node in placeholders]
    examples = _bind_examples(graph_module, args, kwargs)
    normalized = _normalize_dynamic_shapes(dynamic_shapes, names, model, args, kwargs)

    meta = graph_module.meta
    meta["user_signature"] = graph_module.signature
    meta["state_targets"] = dict(tracer.state_targets)
    constants: dict[str, Any] = {}
    for target, (_node, kind, persistent) in tracer.state_targets.items():
        if kind == "constant" or (kind == "buffer" and not persistent):
            constants[target] = _resolve_attribute(model, target)
    meta["constants"] = constants

    mutations = _detect_mutations(graph_module, tracer.state_targets)
    user_out_spec = tree_flatten(graph_module.graph.output_node.args[0])[1]
    ranges, equality_constraints = _apply_dynamic_shape_constraints(
        graph_module, _combine_args(model, args, kwargs), normalized
    )
    _rewrite_container_reads(graph_module, mutations)
    _restructure_output(graph_module, mutations)
    if mutations or ranges:
        graph_module.recompile()
    meta["num_mutations"] = len(mutations)
    meta["out_spec"] = user_out_spec
    meta["in_spec"] = tree_flatten((tuple(args), dict(kwargs)))[1]
    graph_module.signature = _flat_signature(
        graph_module, tracer.state_targets, examples
    )

    meta["module_calls"] = list(tracer.module_calls)
    signature = _graph_signature(graph_module, tracer.state_targets, mutations)
    user_output_args = [
        spec.arg
        for spec in signature.output_specs
        if spec.kind is OutputKind.USER_OUTPUT
    ]
    root_entry = ModuleCallEntry(
        "",
        ModuleCallSignature(
            inputs=[TensorArgument(node.name) for node in placeholders],
            outputs=user_output_args,
            in_spec=meta.get("in_spec"),
            out_spec=meta.get("out_spec"),
            forward_arg_names=[node.name for node in placeholders],
        ),
    )
    call_entries: list[ModuleCallEntry] = []
    for record in tracer.module_calls:
        arg_specs = [
            TensorArgument(value) if isinstance(value, str) else ConstantArgument("", value)
            for value in record["args"]
        ]
        call_entries.append(
            ModuleCallEntry(
                record["fqn"],
                ModuleCallSignature(
                    inputs=arg_specs,
                    outputs=[TensorArgument(name) for name in record["result"]],
                    in_spec=record.get("in_spec"),
                    out_spec=record.get("out_spec"),
                    forward_arg_names=[
                        value if isinstance(value, str) else f"arg_{index}"
                        for index, value in enumerate(record["args"])
                    ],
                ),
            )
        )
    calls = [root_entry, *call_entries]
    program = ExportedProgram(
        graph_module=graph_module,
        graph_signature=signature,
        example_inputs=examples,
        dynamic_shapes=normalized,
        module_call_graph=calls,
        range_constraints=ranges,
        equality_constraints=equality_constraints,
    )
    program.validate()
    return program

[docs]
def export(
    model: Callable[..., Any],
    *args: Any,
    dynamic_shapes: Any = None,
    strict: bool = False,
    preserve_module_call_signature: Any = (),
    **kwargs: Any,
) -> ExportedProgram:
    """Capture a callable and return an executable graph program.

    Args:
        model: an ``nn.Module`` or plain callable; child modules are inlined.
        args/kwargs: example inputs binding argument defaults.
        dynamic_shapes: dimension specification per argument (dict, sequence,
            :class:`ShapesCollection`, or :class:`AdditionalInputs`).
        strict: reserved for callers of the strict capture contract; capture
            validation is identical in both modes.
        preserve_module_call_signature: submodule paths whose call metadata is
            recorded in ``module_call_graph`` for module-level tooling.
    """

    if not callable(model):
        raise TypeError(f"model must be callable, got {type(model).__name__}")
    if isinstance(model, object) and hasattr(model, "named_modules"):
        known = {name for name, _ in model.named_modules(remove_duplicate=True)}
        unknown = [p for p in preserve_module_call_signature if p not in known]
        if unknown:
            raise GraphCaptureError(
                f"preserve_module_call_signature paths {sorted(unknown)} do not "
                f"exist on the model; known: {sorted(known)}"
            )
    program = _capture(model, args, kwargs, dynamic_shapes)
    if preserve_module_call_signature:
        existing = {entry.fqn for entry in program.module_call_graph}
        program.module_call_graph.extend(
            ModuleCallEntry(path)
            for path in preserve_module_call_signature
            if path not in existing
        )
    del strict
    return program

[docs]
def export_for_training(
    model: Callable[..., Any],
    *args: Any,
    dynamic_shapes: Any = None,
    **kwargs: Any,
) -> ExportedProgram:
    """Capture a callable while retaining its mutable training state."""

    return export(model, *args, dynamic_shapes=dynamic_shapes, **kwargs)

def draft_export(
    model: Callable[..., Any],
    *args: Any,
    dynamic_shapes: Any = None,
    **kwargs: Any,
) -> Any:
    """Capture with failure reporting; see :mod:`tensorplay.export._draft_export`."""

    from ._draft_export import draft_export as _draft_export

    return _draft_export(model, *args, dynamic_shapes=dynamic_shapes, **kwargs)
```