yobx.builder.onnxscript.bridge_graph_builder#

class yobx.builder.onnxscript.bridge_graph_builder.OnnxScriptGraphBuilder(target_opset_or_existing_proto: int | Dict[str, int], ir_version: int | None = None)[source]#

Bridge builder that exposes a yobx-compatible API over onnxscript’s IR. It takes onnxscript GraphBuilder implements the API Expected API.

Parameters:

  • target_opset_or_existing_proto – Either a single opset version (int) or a mapping {domain: version} (Dict[str, int]). For example 18 or {"": 18, "com.microsoft": 1}.

  • ir_version – ONNX IR version to use when exporting. When None a sensible default is chosen based on the main opset version.

The builder wraps an onnxscript._internal.builder.GraphBuilder and keeps a name → ir.Value registry so that callers can reference previously created tensors by their string name (as in the yobx API) rather than by ir.Value handle.

Typical usage:

gr = OnnxScriptGraphBuilder(18)
gr.make_tensor_input("X", TensorProto.FLOAT, (None, 4))
w_name = gr.make_initializer("W", np.eye(4, dtype=np.float32))
(y_name,) = gr.make_node("MatMul", ["X", w_name], 1, name="mm")
gr.make_tensor_output(y_name, TensorProto.FLOAT, (None, 4))
proto = gr.to_onnx()
add_domain(domain: str, version: int = 1) → None[source]#

Deprecated. Use set_opset() instead.

get_debug_msg() → str[source]#

Returns information useful for understanding where an error could come from.

Returns:

a summary of the current graph state as a string

get_device(name: str) → int[source]#

Returns the device. This is not supported right now.

get_opset(domain: str, exc: bool = True) → int[source]#

Returns the opset version for a specific domain.

Parameters:

  • domain – domain name

  • exc – raise an exception if missing

Returns:

version

get_shape(name: str) → DYNAMIC_SHAPE[source]#

Returns the shape.

get_type(name: str) → int[source]#

Returns the type.

get_value(name: str) → Value[source]#

Returns the ir.Value associated with name.

Parameters:

name – Tensor name.

Raises:

KeyError – When name has not been registered.

has_device(name: str) → int[source]#

Tells if a value has a device. This is not supported right now.

has_name(name: str) → bool[source]#

Returns True when name is a known value in this graph.

Parameters:

name – Tensor name to query.

has_opset(domain: str) → int[source]#

Returns the opset version for a domain, or 0 if the domain is not registered.

Parameters:

domain – domain name

Returns:

opset version, or 0 if the domain is unknown

has_shape(name: str) → bool[source]#

Tells if a value has a shape.

has_type(name: str) → int[source]#

Tells if a value has a type.

property inner_builder#

The underlying onnxscript._internal.builder.GraphBuilder. Use this to access onnxscript-native functionality that is not exposed through the yobx-compatible bridge API.

Returns:

onnxscript._internal.builder.GraphBuilder

property input_names: List[str]#

Returns input names.

property main_opset: int#

Returns the opset for the main domain (assuming it is used).

make_initializer(name: str, value: Any) → str[source]#

Adds an initializer tensor and return its name.

Parameters:

  • name – Name for the initializer. May be an empty string "" in which case a unique name is generated automatically.

  • value – Initializer data. Supported types: numpy.ndarray, scalar int/float, onnx.TensorProto.

Returns:

The final registered name (may differ from name when name is empty).

make_node(op_type: str, inputs: str | List[str], outputs: int | str | List[str] = 1, domain: str = '', attributes: List[AttributeProto] | None = None, name: str | None = None, **kwargs: Any) → str | Tuple[str, ...][source]#

Creates an ONNX node and return its output name(s).

Parameters:

  • op_type – ONNX operator type (e.g. "Relu", "MatMul").

  • inputs – Input tensor name(s). Each name must have been registered previously via make_tensor_input(), make_initializer(), or a previous make_node() call. Pass an empty string "" for optional absent inputs.

  • outputs – Either an int (number of outputs, names are auto-generated), a single output str, or a list of output str names.

  • domain – Operator domain (default: "" = standard ONNX).

  • attributes – Additional onnx.AttributeProto instances to attach to the node (supplementary to kwargs).

  • name – Optional node name (for debugging / profiling).

  • kwargs – Operator attributes as Python primitives (int, float, str, List[int], …).

Returns:

The output name when a single output is created, or a list of names when multiple outputs are created.

make_tensor_input(name: str, elem_type: int | None = None, shape: Sequence[int | str | None] | None = None, device: int | None = None) → str[source]#

Adds a graph input and return its name.

Parameters:

  • name – Input tensor name.

  • elem_type – ONNX element type (e.g. TensorProto.FLOAT). Pass None or 0 if unknown (only valid for function graphs).

  • shape – Tensor shape. Use None for a fully unknown dimension and a str for a symbolic / dynamic dimension.

  • device – unused for the time being

Returns:

The registered name (same as name).

make_tensor_output(name: str | List[str], elem_type: int | None = None, shape: Sequence[int | str | None] | None = None, indexed: bool = False, allow_untyped_output: bool = False) → str | List[str][source]#

Registers an existing value as a graph output and return its name.

Parameters:

  • name – Name (or list of names) of the tensor(s) to mark as graph output(s). Must already exist in this builder (i.e. have been created by make_tensor_input(), make_initializer(), or make_node()).

  • elem_type – Optional element type hint; used to set the type on the ir.Value if it was not already inferred.

  • shape – Optional shape hint; used to set / refine the shape on the ir.Value if not already inferred.

  • indexed – unused

  • allow_untyped_output – allows output with no shape and/or no type

Returns:

The name (or list of names), matching the name argument.

onnx_dtype_to_np_dtype(itype: int) → dtype[source]#

See yobx.helpers.onnx_helper.tensor_dtype_to_np_dtype().

property op: OpsetProtocol#

Returns the shortcut to OpsetProtocal.

property output_names: List[str]#

Returns output names.

set_opset(domain: str, version: int = 1) → None[source]#

Sets the opset version for a domain. Checks the version is the same if it already exists.

Parameters:

  • domain – domain name to register

  • version – opset version for the domain

set_shape(name: str, shape: DYNAMIC_SHAPE, allow_zero: bool = False)[source]#

Sets the shape.

set_type(name: str, itype: int)[source]#

Sets the type.

set_type_shape_unary_op(name: str, input_name: str, itype: int | None = None) → bool[source]#

Propagates type and shape from input_name to name for a unary op.

This is a convenience helper used when emitting elementwise unary operators (Abs, Exp, Relu, etc.) where the output has the same type and shape as the input.

Parameters:

  • name – output tensor name whose type/shape should be set

  • input_name – input tensor name to copy type/shape from

  • itype – override element type; if None the input’s element type is used

Returns:

True when shape information was available and set, None/falsy when it could not be determined

to_onnx(large_model: bool = False, external_threshold: int = 1024, inline: bool = True) → ExportArtifact[source]#

Exports the graph as an ONNX ModelProto.

Parameters:

  • large_model – if True returns a onnx.model_container.ModelContainer, it lets the user to decide later if the weights should be part of the model or saved as external weights

  • external_threshold – if large_model is True, every tensor above this limit (in bytes) is stored as external

  • inline – inline local function it any (this is currently not used)

Returns:

artifact

unique_name(prefix: str) → str[source]#

Returns a unique name.

class yobx.builder.onnxscript.bridge_graph_builder.OnnxScriptGraphBuilderOpset(builder)[source]#

Implements yobx.typing.OpsetProtocol.

yobx.builder.onnxscript.bridge_graph_builder.to_ir_dtype(elem_type: int | None) → DataType | None[source]#

Converts an ONNX TensorProto element-type integer to ir.DataType.

Parameters:

elem_type – ONNX element type (e.g. TensorProto.FLOAT == 1), or None / 0 for unknown.

Returns:

Corresponding ir.DataType, or None when unknown.

yobx.builder.onnxscript.bridge_graph_builder.to_ir_shape(shape: Sequence[int | str | None] | None) → Shape | None[source]#

Converts a yobx-style shape tuple to ir.Shape.

Parameters:

shape – A sequence of dimension sizes. Each element may be an int (static), a str (symbolic / dynamic), or None (fully unknown dimension).

Returns:

ir.Shape, or None when shape itself is None.

yobx.builder.onnxscript.bridge_graph_builder.value_to_ir_tensor(value: Any, name: str) → TensorProtocol[source]#

Converts an initializer value to an ir.TensorProtocol.

Supported input types:

  • numpy.ndarray

  • scalar Python int or float (promoted to 0-D arrays)

  • onnx.TensorProto (converted via ir.from_proto())

  • Any object already satisfying ir.TensorProtocol

Parameters:

  • value – The raw initializer value.

  • name – Name that should be associated with the resulting tensor.

Returns:

An ir.TensorProtocol suitable for passing to onnxscript._internal.builder.GraphBuilder.initializer().

Raises:

TypeError – When value has an unsupported type.