KeyError: (0, 0) in ONNXSurrogate.populate_block due to key format mismatch between IDAES and OMLT #1733
Description
Description:
I am encountering a KeyError: (0, 0) when attempting to build an IDAES SurrogateBlock using an ONNXSurrogate object created from a PyTorch-exported ONNX model.
It appears there is a mismatch between how IDAES handles input/output indices and how the latest version of OMLT identifies nodes in the ONNX graph.
Environment:
IDAES Version: 2.9.0
OMLT Version: 1.2.2
PyTorch Version: 2.9.1
Operating System: Windows 11
Steps to Reproduce:
Train a simple ANN model in PyTorch.
Export the model to ONNX format using torch.onnx.export.
Create an IDAES ONNXSurrogate object with input_bounds and OffsetScaler.
Attempt to build the model via m.fs.surrogate.build_model(surrogate_obj, ...).
Error Traceback:
Python
Traceback (most recent call last):
File "...", line 95, in
m.fs.surrogate.build_model(surrogate_obj, input_vars=inputs, output_vars=outputs)
File ".../idaes/core/surrogate/surrogate_block.py", line 126, in build_model
surrogate_object.populate_block(self, additional_options=kwargs)
File ".../idaes/core/surrogate/onnx_surrogate.py", line 152, in populate_block
self.populate_block_with_net(block, formulation_object)
File ".../idaes/core/surrogate/omlt_base_surrogate_class.py", line 163, in populate_block_with_net
block.nn.build_formulation(formulation_object)
...
File ".../omlt/formulation.py", line 92, in
k: (float(scaled_input_bounds[k][0]), float(scaled_input_bounds[k][1]))
KeyError: (0, 0)
Proposed Cause:
In idaes.core.surrogate.onnx_surrogate (and the base OMLT surrogate class), the code generates input bounds using integer keys:
Python
idaes/core/surrogate/onnx_surrogate.py
input_bounds = dict(enumerate(self.input_bounds().values()))
This results in a dictionary like {0: (lb, ub), 1: (lb, ub)}.
However, OMLT's load_onnx_neural_network and its formulation builders now seem to identify input nodes using tuples (GraphID, NodeID), such as (0, 0). When OMLT attempts to look up the scaled bounds using the tuple key (0, 0), it fails because IDAES only provided integer keys.
Additionally, I noticed that NetworkDefinition attributes seem inconsistent across versions (e.g., input_indexes vs input_nodes), which complicates manual workarounds.
Expected Behavior:
ONNXSurrogate should dynamically detect the input/output node keys from the OMLT NetworkDefinition to ensure compatibility regardless of whether the keys are integers or tuples.