onnx_diagnostic.helpers.torch_fx_graph_helper

onnx_diagnostic.helpers.torch_fx_graph_helper.prepare_args_kwargs(torch_results: Dict[str, Any], node: Node) → Tuple[Tuple[Any, ...], Dict[str, Any]]

Prepares args and kwargs before executing a fx node.

Parameters:

• torch_results – existing results

• node – node to execute

Returns:

new args and kwargs

onnx_diagnostic.helpers.torch_fx_graph_helper.run_fx_node(node: Node, args: Tuple[Any, ...], kwargs: Dict[str, Any] | None = None) → Tuple[Any, ...]

Executes a node

Parameters:

• node – runs a node

• args – unnamed inputs to the node

• kwargs – named inputs to the node

Returns:

results

onnx_diagnostic.helpers.torch_fx_graph_helper.validate_fx_outputs(node: Node, outputs: Tuple[Any, ...]) → None

Validates the outputs of a node using metadata stored in the node.

Parameters:

• node – node

• outputs – outputs

onnx_diagnostic.helpers.torch_fx_graph_helper.validate_fx_tensor(node: Node, tensor: Tensor, expected_shape: Tuple[Any, ...]) → None

Validates the shape of tensor is expected.

Parameters:

• node – node

• tensor – tensor

• expected_shape – expected shape