onnx_light.tools#

Standalone helpers built on top of onnx_light.

This sub-package collects small utilities that are not part of the upstream onnx API but are useful when working with ONNX models through onnx_light.

Currently provided helpers:

onnx_light.tools.build_kernel_inventory(*, library: str = 'onnx_light', dispatch_table_path: Path | str | None = None, kernel_source_root: Path | str | None = None, tuning_report: dict[str, Any] | None = None) list[dict[str, Any]]#

Builds the full kernel-path inventory (Step D).

Parameters:
  • library – Tuning library identifier passed to onnx_light.kernel_tuning.kernel_tuning_parameters() when tuning_report is not supplied.

  • dispatch_table_path – Optional override of the dispatch-table source path.

  • kernel_source_root – Optional override of the kernel implementation source root.

  • tuning_report – Optional pre-fetched kernel_tuning_parameters() result, mainly for tests that must not depend on the compiled extension.

Returns:

A list of inventory rows, sorted by (domain, op_type, element_type, implementation). Every registered kernel path appears exactly once.

onnx_light.tools.compute_value_and_node_tags(graph_or_nodes_or_function: Any, verbose: int = 0) tuple[dict[str, str], list[str]]#

Infers semantic shape/axes/weight/ambiguous tags for values and nodes.

Returns:

A pair (value_tags, node_tags) where value_tags maps value names to tags and node_tags is ordered like the processed node list.

onnx_light.tools.get_cpu_descriptor() dict[str, Any]#

Returns a best-effort, portable CPU descriptor for the local machine.

Missing information is omitted rather than replaced by an invented value.

Returns:

A mapping with at least architecture and logical_cores.

onnx_light.tools.infer_value_and_node_tags(graph_or_nodes_or_function: Any, verbose: int = 0) tuple[dict[str, str], list[str]]#

Infers semantic shape/axes/weight/ambiguous tags for values and nodes.

Returns:

A pair (value_tags, node_tags) where value_tags maps value names to tags and node_tags is ordered like the processed node list.

onnx_light.tools.pretty_onnx(onx: Any, with_attributes: bool = False, highlight: set[str] | None = None, shape_inference: bool = False, include_node_tags: bool = False, include_inplace: bool = False, include_release: bool = False) str#

Returns a compact, human-readable string for any ONNX proto.

The argument may be a ModelProto, GraphProto, FunctionProto, NodeProto, ValueInfoProto, TypeProto, AttributeProto, TensorProto or a file path (str). The rendering style mirrors yobx.helpers.onnx_helper.pretty_onnx: nodes appear as OpType(inputs) -> outputs, value infos as dtype[shape] name, and graphs/models as a list of opsets, inputs, initializers, nodes and outputs.

Parameters:
  • onx – ONNX proto, or a path to a model file.

  • with_attributes – when True, node attributes are appended to the node line (each name=value pair on its own line, or inline after --- when there is a single attribute).

  • highlight – optional set of tensor names to wrap in ** markers in the rendered I/O lists.

  • shape_inference – when True and onx is a model, run onnx_light shape inference before rendering.

  • include_node_tags – when True, nodes that carry a onnx_light.node_tag metadata entry (shape, axes or weight) are prefixed with [tag] in the rendered output.

  • include_inplace – when True, nodes that carry onnx_light.inplace_reuse metadata have the inplace reuse opportunities appended to their line, e.g. inplace: out0=in0(equal).

  • include_release – when True, nodes that carry onnx_light.release_after metadata have the release hints appended to their line, and nodes with onnx_light.not_used_after metadata include not used after: ....

Returns:

the formatted text.

onnx_light.tools.run_benchmark_corpus(*, cases: tuple[dict[str, Any], ...] | None = None, cpu_policies: tuple[tuple[str, int], ...] | None = None, repeat: int = 5, warmup: int = 2, seed: int = 0, collect_diagnostics: bool = True) list[dict[str, Any]]#

Runs every (case, shape, cpu policy) combination in the corpus.

Parameters:
  • cases – Benchmark cases; defaults to BENCHMARK_CORPUS.

  • cpu_policies(label, num_threads) pairs; defaults to CPU_POLICIES.

Returns:

A flat list of result rows; see run_kernel_baseline_report() for the combined report schema.

onnx_light.tools.run_kernel_baseline_report(*, cases: tuple[dict[str, Any], ...] | None = None, cpu_policies: tuple[tuple[str, int], ...] | None = None, repeat: int = 5, warmup: int = 2, seed: int = 0, collect_diagnostics: bool = True) dict[str, Any]#

Produces the combined Step D + Step E machine-readable report.

This is the single entry point behind python -m onnx_light kernel-baseline: it does not modify the kernel tuning cache and does not invoke onnxruntime, so only native onnx-light kernel execution enters the report.

Parameters:
  • cases – Benchmark cases; defaults to BENCHMARK_CORPUS.

  • cpu_policies(label, num_threads) pairs; defaults to CPU_POLICIES.

Returns:

A mapping with cpu_descriptor, inventory (Step D rows) and benchmarks (Step E rows).

onnx_light.tools.to_dot(model_or_graph: Any, *, direction: str = 'TB', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False) str#

Renders an ONNX ModelProto or GraphProto as a Graphviz DOT string.

Parameters:
  • model_or_graph – A ModelProto or GraphProto instance. Both onnx_light and onnx messages are accepted.

  • direction – Graphviz layout direction; one of "TB" (top-to-bottom), "BT" (bottom-to-top), "LR" (left-to-right), or "RL" (right-to-left). Defaults to "TB".

  • include_initializers – When True, initializers are rendered as separate (cylinder) nodes and connected to the consumers they feed. When False, initializer tensors are not shown.

  • include_shapes – When True, tensor type/shape information available in graph inputs, outputs, value_info and initializers is appended to the corresponding node labels.

  • include_attributes – When True, node attribute names are listed inside the operator label.

  • include_inplace – When True, the in-place reuse opportunities recorded in each node’s metadata_props (under the onnx_light.inplace_reuse key) are appended to the operator label, for example inplace: out0=in1(equal).

Returns:

The DOT source as a single str (newline-separated). Write it to a .dot file and render it with, for example, dot -Tsvg model.dot -o model.svg.

Raises:
  • TypeError – If model_or_graph is neither a ModelProto nor a GraphProto.

  • ValueError – If direction is not a supported Graphviz direction.

The example below builds a small Abs chain and renders the DOT source:

    digraph onnx {
        rankdir="TB";
        node [fontname="Helvetica", fontsize=10];
        edge [fontname="Helvetica", fontsize=9];
        "t_X" [label="X\nfloat[3,4]", shape=ellipse, style="filled", fillcolor="#cde4ff", color="#3a6ea5"];
        "n_Abs_0" [label="Abs", shape=box, style="filled", fillcolor="#d4ecd4", color="#3a8c3a"];
        "t_X" -> "n_Abs_0" [label="float[3,4]"];
        "n_Abs_0" -> "t_A";
        "n_Abs_1" [label="Abs", shape=box, style="filled", fillcolor="#d4ecd4", color="#3a8c3a"];
        "t_A" -> "n_Abs_1";
        "n_Abs_1" -> "t_B";
        "n_Abs_2" [label="Abs", shape=box, style="filled", fillcolor="#d4ecd4", color="#3a8c3a"];
        "t_B" -> "n_Abs_2";
        "n_Abs_2" -> "t_Y" [label="float[3,4]"];
        "t_Y" [label="Y\nfloat[3,4]", shape=ellipse, style="filled", fillcolor="#ffe1b3", color="#a35a00"];
    }
onnx_light.tools.to_dot_graph(graph: Any, *, direction: str = 'TB', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False) str#

Renders a GraphProto as a Graphviz DOT string.

See to_dot() for the meaning of every parameter.

onnx_light.tools.to_mermaid(model_or_graph: Any, *, direction: str = 'TB', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False, include_release: bool = False) str#

Render an ONNX ModelProto or GraphProto as a Mermaid flowchart.

Parameters:
  • model_or_graph – A ModelProto or GraphProto instance. Both onnx_light and onnx messages are accepted.

  • direction – Mermaid flowchart direction; one of "TB", "TD", "BT", "LR" or "RL". Defaults to "TB" (top-to-bottom).

  • include_initializers – When True, initializers are rendered as separate (cylinder) nodes and connected to the consumers they feed. When False, initializer tensors are not shown.

  • include_shapes – When True, tensor type/shape information available in graph inputs, outputs, value_info and initializers is appended to the corresponding node labels.

  • include_attributes – When True, node attribute names are listed inside the operator label.

  • include_inplace – When True, the in-place reuse opportunities recorded in each node’s metadata_props (under the onnx_light.inplace_reuse key) are appended to the operator label, for example inplace: out0=in1(equal).

  • include_release – When True, the post-execution release hints recorded in each node’s metadata_props (under the onnx_light.release_after key, and optional last-use hints under onnx_light.not_used_after) are appended to the operator label, for example release: A, B; not used after: X.

Returns:

The Mermaid source as a single str (newline-separated). The returned text is not wrapped in a fenced code block so the caller can choose between "```mermaid\n...\n`”`` for Markdown and the .. mermaid:: directive for Sphinx.

Raises:
  • TypeError – If model_or_graph is neither a ModelProto nor a GraphProto.

  • ValueError – If direction is not a supported Mermaid flowchart direction.

The example below builds a small Abs chain, runs shape inference and records the in-place reuse opportunities into the graph metadata with onnx_light.onnx_core.shape_inference.write_inplace_reuse_to_metadata(), then renders the annotated flowchart with include_inplace=True:

        flowchart TB
    t_X(["X<br/>float[3,4]"]):::onnxInput
    n_Abs_0["Abs"]:::onnxOp
    t_X -->|"float[3,4]"| n_Abs_0
    n_Abs_0 --> t_A
    n_Abs_1["Abs<br/>inplace: out0=in0(equal)"]:::onnxOp
    t_A --> n_Abs_1
    n_Abs_1 --> t_B
    n_Abs_2["Abs<br/>inplace: out0=in0(equal)"]:::onnxOp
    t_B --> n_Abs_2
    n_Abs_2 -->|"float[3,4]"| t_Y
    t_Y(["Y<br/>float[3,4]"]):::onnxOutput
    classDef onnxInput fill:#cde4ff,stroke:#3a6ea5,color:#000;
    classDef onnxOutput fill:#ffe1b3,stroke:#a35a00,color:#000;
    classDef onnxInitializer fill:#eeeeee,stroke:#888,color:#000;
    classDef onnxOp fill:#d4ecd4,stroke:#3a8c3a,color:#000;
    classDef onnxTagShape fill:#f4d6ff,stroke:#8744a2,color:#000;
    classDef onnxTagAxes fill:#ffe9a8,stroke:#9e7a00,color:#000;
    classDef onnxTagWeight fill:#e0e0e0,stroke:#666666,color:#000;
    classDef onnxTagAmbiguous fill:#ffd9d9,stroke:#a33a3a,color:#000;
    
onnx_light.tools.to_mermaid_graph(graph: Any, *, direction: str = 'TB', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False, include_release: bool = False) str#

Render a GraphProto as a Mermaid flowchart.

See to_mermaid() for the meaning of every parameter.

onnx_light.tools.to_svg(model_or_graph: Any, *, direction: str = 'TB', layout: str = 'layered', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False, include_release: bool = False) str#

Renders an ONNX ModelProto or GraphProto as an SVG image.

Parameters:
  • model_or_graph – A ModelProto or GraphProto instance. Both onnx_light and onnx messages are accepted.

  • direction – Layout direction; "TB" (or "TD") lays the graph out top-to-bottom and "LR" left-to-right. Defaults to "TB".

  • layout – Positioning strategy for the boxes. "layered" (the default) uses the built-in layered layout with barycenter crossing reduction. "umap" instead derives the box positions from a two-dimensional UMAP embedding of the graph’s connectivity, which requires the optional umap-learn package to be installed.

  • include_initializers – When True, initializers are rendered as separate (dashed) boxes connected to their consumers. When False, initializer tensors are not shown.

  • include_shapes – When True, tensor type/shape information available in graph inputs, outputs, value_info and initializers is appended to the corresponding box labels.

  • include_attributes – When True, node attribute names are listed inside the operator label.

  • include_inplace – When True, the in-place reuse opportunities recorded in each node’s metadata_props (under the onnx_light.inplace_reuse key) are appended to the operator label, for example inplace: out0=in1(equal).

  • include_release – When True, the post-execution release hints recorded in each node’s metadata_props (under the onnx_light.release_after key, and optional last-use hints under onnx_light.not_used_after) are appended to the operator label, for example release: A, B; not used after: X.

Returns:

A self-contained SVG document as a single str.

Raises:
  • TypeError – If model_or_graph is neither a ModelProto nor a GraphProto.

  • ValueError – If direction or layout is not supported.

The example below builds a small Abs chain, runs shape inference and records the in-place reuse opportunities into the graph metadata with onnx_light.onnx_core.shape_inference.write_inplace_reuse_to_metadata(), then renders the annotated diagram with include_inplace=True:

X · float[3,4] A B Y · float[3,4] Xfloat[3,4] Abs Absinplace: out0=in0(equal) Absinplace: out0=in0(equal) float[3,4]
onnx_light.tools.to_svg_graph(graph: Any, *, direction: str = 'TB', layout: str = 'layered', include_initializers: bool = True, include_shapes: bool = True, include_attributes: bool = False, include_inplace: bool = False, include_release: bool = False) str#

Renders a GraphProto as an SVG image.

See to_svg() for the meaning of every parameter.

onnx_light.tools.translate(proto: Any, api: str = 'onnx-compact') str#

Translates an ONNX model or graph into Python code that rebuilds it.

Parameters:
  • proto – a ModelProto or GraphProto (or a file path to load).

  • api – target flavour, "onnx-compact" (default) or "builder".

Returns:

The generated Python code as a string (without the import header, see translate_header()).

onnx_light.tools.translate_header(api: str = 'onnx-compact') str#

Returns the import header required by the code produced by translate().

Parameters:

api – target flavour, "onnx-compact" or "builder".

Returns:

The import header as a string ending with a trailing newline.

onnx_light.tools.write_value_and_node_tags_to_metadata(graph_or_nodes_or_function: Any) None#

Writes inferred shape/axes/weight/ambiguous tags into metadata.

Returns:

None. The function mutates metadata fields in place.