utils

Utility functions for AutoCast.

This module provides common utility functions used across the AutoCast package. It includes functions for graph traversal, tensor type inference, model validation, and mapping setup between nodes, initializers, and value info. These utilities support the core functionality of model precision conversion.

Functions

`get_cast_to_type`	Get the target type from a Cast node.
`get_consumer_nodes`	Get all consumer nodes for a given tensor name.
`get_producer_nodes`	Get all producer nodes for a given tensor name.
`get_unique_consumer_node`	Get a single consumer node and raise exception if there are multiple consumers.
`setup_mappings`	Setup and return mappings for model components.

get_cast_to_type(cast_node)

Get the target type from a Cast node.

Parameters:: cast_node (NodeProto) – The Cast node to extract type from.
Returns:: The target type value from the Cast node’s ‘to’ attribute.
Return type:: int
Raises:: ValueError – If the Cast node does not have a ‘to’ attribute.

get_consumer_nodes(model, tensor_name)

Get all consumer nodes for a given tensor name.

Parameters:

model (ModelProto) – The ONNX model to search.
tensor_name (str) – Name of the tensor to find consumers for.

Returns:

List of nodes that consume the tensor.

Return type:

list[onnx.NodeProto]

get_producer_nodes(model, tensor_name)

Get all producer nodes for a given tensor name.

Parameters:

model (ModelProto) – The ONNX model to search.
tensor_name (str) – Name of the tensor to find producers for.

Returns:

List of nodes that produce the tensor.

Return type:

list[onnx.NodeProto]

get_unique_consumer_node(model, tensor_name)

Get a single consumer node and raise exception if there are multiple consumers.

Parameters:

model (ModelProto) – The ONNX model to search.
tensor_name (str) – Name of the tensor to find consumer for.

Returns:

The single consumer node.

Return type:

onnx.NodeProto

Raises:

Exception – If there is not exactly one consumer node.

setup_mappings(model)

Setup and return mappings for model components.

Parameters:

model (ModelProto) – ONNX model to create mappings for.

Returns:

value_info_map: Mapping of names to value infos.
initializer_map: Mapping of names to initializers.
node_to_init_map: Mapping of node names to their initializer inputs.

Return type:

Tuple containing