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.

prefix_name_context(prefix: str) → Generator[source]#

Context manager that pushes prefix onto the prefix stack.

While active, unique_name() prepends the joined stack to every requested name so that all tensor names produced inside the block are scoped to the current context. The prefix is removed when the block exits, even if an exception is raised.

Parameters:

prefix – scope prefix to push (e.g. a pipeline step name)

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.