_tape.py

# Copyright (c) ONNX Project Contributors
# SPDX-License-Identifier: Apache-2.0
"""Convenience methods for constructing the IR."""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import (
    Any,
    Optional,
)

import onnx_ir as ir
from onnx_ir import _convenience

# A type representing the domains/versions used in creating nodes in IR.
UsedOpsets = set[tuple[str, Optional[int]]]


class Tape:
    """Tape class.

    A tape is a recorder that collects nodes and initializers that are created so
    that they can be used for creating a graph.

    Example::

        import onnx_ir as ir

        tape = ir.tape.Tape()
        a = tape.initializer(ir.tensor([1, 2, 3], name="a"))
        b: ir.Value = ...
        c: ir.Value = ...
        x = tape.op("Add", [a, b], attributes={"alpha": 1.0})
        y = tape.op("Mul", [x, c], attributes={"beta": 2.0})
        model = ir.Model(
            graph := ir.Graph(
                inputs=[b, c],
                outputs=[y],
                nodes=tape.nodes,
                initializers=tape.initializers
                opset_imports={"": 20},
            ),
            ir_version=10,
        )

    Attributes:
        graph_like: The graph to append the new nodes and initializers to. When
            it is None, the nodes and initializers are creating without owned by a graph.
            Initializers will not be added to functions because it is not supported by ONNX.
    """

    def __init__(self, graph_like: ir.Graph | ir.Function | None = None) -> None:
        self._nodes: list[ir.Node] = []
        self._initializers: list[ir.Value] = []
        self._used_opsets: UsedOpsets = set()
        self.graph_like = graph_like

    def __repr__(self) -> str:
        return f"Tape(nodes={self._nodes}, initializers={self._initializers})"

    @property
    def nodes(self) -> Sequence[ir.Node]:
        return tuple(self._nodes)

    @property
    def initializers(self) -> Sequence[ir.Value]:
        return tuple(self._initializers)

    @property
    def used_opsets(self) -> UsedOpsets:
        return self._used_opsets

    def op(
        self,
        op_type: str,
        inputs: Sequence[ir.Value | None],
        attributes: Mapping[str, _convenience.SupportedAttrTypes] | None = None,
        *,
        domain: str = "",
        overload: str = "",
        version: int | None = None,
        graph: ir.Graph | None = None,
        name: str | None = None,
        doc_string: str | None = None,
        metadata_props: dict[str, str] | None = None,
        output: ir.Value | None = None,
    ) -> ir.Value:
        if attributes is None:
            attrs: Sequence[ir.Attr] = ()
        else:
            attrs = _convenience.convert_attributes(attributes)
        output_kwargs: dict[str, Any]
        if output is None:
            output_kwargs = dict(num_outputs=1)
        else:
            output_kwargs = dict(outputs=[output])
        node = ir.Node(
            domain,
            op_type,
            inputs,
            attributes=attrs,
            **output_kwargs,
            overload=overload,
            version=version,
            graph=graph or self.graph_like,
            name=name,
            doc_string=doc_string,
            metadata_props=metadata_props,
        )
        self._nodes.append(node)
        self._used_opsets.add((domain, version))

        return node.outputs[0]

    def op_multi_out(
        self,
        op_type: str,
        inputs: Sequence[ir.Value | None],
        attributes: Mapping[str, _convenience.SupportedAttrTypes] | None = None,
        *,
        num_outputs: int | None = None,
        outputs: Sequence[ir.Value] | None = None,
        domain: str = "",
        overload: str = "",
        version: int | None = None,
        graph: ir.Graph | None = None,
        name: str | None = None,
        doc_string: str | None = None,
        metadata_props: dict[str, str] | None = None,
    ) -> Sequence[ir.Value]:
        if num_outputs is None and outputs is None:
            raise ValueError("Either num_outputs or outputs must be provided.")
        if num_outputs is not None and outputs is not None:
            raise ValueError("Both num_outputs and outputs cannot be provided simultaneously.")
        output_kwargs: dict[str, Any]
        if outputs is None:
            output_kwargs = dict(num_outputs=num_outputs)
        else:
            output_kwargs = dict(outputs=outputs)
        if attributes is None:
            attrs: Sequence[ir.Attr] = ()
        else:
            attrs = _convenience.convert_attributes(attributes)
        node = ir.Node(
            domain,
            op_type,
            inputs,
            attributes=attrs,
            **output_kwargs,
            overload=overload,
            version=version,
            graph=graph or self.graph_like,
            name=name,
            doc_string=doc_string,
            metadata_props=metadata_props,
        )
        self._nodes.append(node)
        self._used_opsets.add((domain, version))

        return node.outputs

    def initializer(self, tensor: ir.TensorProtocol, name: str | None = None) -> ir.Value:
        name = name or tensor.name
        if name is None:
            raise ValueError("Name must be provided for initializer.")
        shape = ir.Shape((d if isinstance(d, int) else d.value) for d in tensor.shape.dims)
        value = ir.Value(
            name=name, shape=shape, type=ir.TensorType(tensor.dtype), const_value=tensor
        )
        self._initializers.append(value)
        if isinstance(self.graph_like, ir.Graph):
            self.graph_like.register_initializer(value)
        return value


class Builder(Tape):
    """An extension of the tape that provides a more convenient API for constructing the IR."""

    def __getattr__(self, op_type: str) -> Any:
        return lambda *args, **kwargs: self._make_node(op_type, args, kwargs)

    def _make_node(self, op_type: str, inputs: Sequence[ir.Value], kwargs: dict[str, Any]):
        domain = kwargs.pop("_domain", "")
        version = kwargs.pop("_version", None)
        outputs = kwargs.pop("_outputs", 1)
        if isinstance(outputs, Sequence):
            num_outputs = len(outputs)
        else:
            assert isinstance(outputs, int)
            num_outputs = outputs

        if num_outputs == 1:
            value = super().op(
                op_type, inputs=inputs, attributes=kwargs, domain=domain, version=version
            )
            if isinstance(outputs, Sequence):
                value.name = outputs[0]
            return value
        values = super().op_multi_out(
            op_type,
            inputs=inputs,
            attributes=kwargs,
            domain=domain,
            version=version,
            num_outputs=num_outputs,
        )
        if isinstance(outputs, Sequence):
            for value, name in zip(values, outputs):
                value.name = name
        return values