ONNXRuntime Error: Expected type

[burakssen]

So the error that I am currently getting on runtime. It fails while it tries to load.

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from 01_couette_flow/couette.onnx failed:Type Error: Type (tensor(float)) of output arg (var_294657) of node (concat_4373) does not match expected type (tensor(int32)).

I am using this model.

couette.onnx.zip

Unfortunatelly it's a pretty huge model around 80 mb. I had to wait for this around 5 hours to get exported from jax2onnx 😅.

[enpasos]

Thanks for the report! The error

Type Error: Type (tensor(float)) of output arg (var_294657) of node (concat_4373) does not match expected type (tensor(int32)).

indicates that a Concat node is producing float, while the downstream graph expects int32—typically a dtype-propagation issue (for example, mixing index tensors with float tensors or missing/extra Cast operations along a shape/index path).

This might be related to #52, where we’re working on fixing incorrect dtype/shape metadata at function boundaries. If it’s the same root cause, it should be resolved once that fix is in—though #52 hasn’t been merged yet.

If you suspect it’s unrelated to #52, or if you just want to be sure, it would be extremely helpful if you could make the ONNX file generation process reproducible. Having the original JAX function and export setup makes it much easier to pinpoint and fix the issue.

[burakssen]

I am not really sure, I don't think the onnx file generation for our model'll help. It takes too much time to get the onnx file, like at least 6 hours or something like that. If you still want it, I'll try to create something minimal but it'll probably be a similar thing.

[enpasos]

Understood — in that case, it’s worth stressing that with any jax2onnx export we aim for at least:

1. Export: The ONNX file is written successfully.
2. Load: The model loads in ONNX Runtime without type/shape errors.
3. Match: The ONNX model is numerically consistent with the original JAX function on representative inputs.

Our plugin/examples suite follows this staircase: we start from small cases that satisfy (1)–(3), then scale up. Given current gaps, this “small → larger → full” progression is the most practical way to catch regressions early. Starting from a 5-hour export that fails at step (2) is rarely tractable; even #52, which is much smaller, shows how tricky dtype/shape issues can be at function boundaries.

Proposal to move forward efficiently:

• Extract or mock the smallest block that still reproduces the load-time dtype mismatch.
• Validate it against (1)–(3). If it fails, we can turn it into a unit test and fix the responsible plugin/boundary.
• Once the minimal case passes, step up in size until the original model path succeeds.

If sharing a full repro isn’t feasible, a minimal synthetic that still fails at (2) is enough for us to harden the boundary and prevent this class of error in the future.

[burakssen]

Ok thank you, I also I think I found a way to decrease the size, we were doing a lot of inner step calculations which I guess created each step continuously, until it ended. I have created a simple one with 1 step it's around 1.6 mb's and I still get the same kind of error.

double_mach_reflection.onnx.zip

[enpasos]

Thanks for the reduced model. I’ve landed fixes for the minimal repros you shared (dtype propagation around Concat/Gather). Can you re-export on the latest main and confirm whether the load error still occurs?

[burakssen]

Thank you,
This is what I get on runtime:

2025-08-22 11:05:26.905605 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'fallback_updates_old_value_0_expanded_for_FILL_OR_DROP_UpdatesReplace_Y_0' source:{-1,-1,-1,28,-1} target:{5,12,28,1,1}. Falling back to lenient merge.
2025-08-22 11:05:26.906360 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_259' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.906484 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_247' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.906656 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_233' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.906830 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_217' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907001 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_201' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907320 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_409' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907435 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_397' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907604 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_383' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907773 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_367' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.907941 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_351' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.918389 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2217' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.918526 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2205' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.918712 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2191' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.918886 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2175' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.919063 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2159' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.919404 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2367' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.919541 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2355' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.919731 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2341' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.919903 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2325' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.920081 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2309' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.931144 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4180' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.931281 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4168' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.931471 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4154' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.931648 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4138' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.931860 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4122' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.933710 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4330' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.933863 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4318' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.934074 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4304' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.934281 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4288' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.934472 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_4272' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.945663 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6143' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.945811 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6131' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.946001 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6117' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.946202 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6101' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.946392 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6085' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.946721 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6293' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.946854 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6281' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.947055 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6267' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.947273 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6251' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:26.947483 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_6235' source:{4,1} target:{1,4,4,1}. Falling back to lenient merge.
2025-08-22 11:05:27.024286 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'safe_updates_0' source:{5,12,28,28,-1} target:{5,12,28,1,1}. Falling back to lenient merge.
2025-08-22 11:05:27.024783 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'shape_rev0_9' source:{2} target:{4}. Falling back to lenient merge.
Traceback (most recent call last):
  File "/Users/burakssen/dev/jaxfluids/jax2onnx-examples/src/2d_examples/load.py", line 12, in <module>
    session = ort.InferenceSession(model_path, session_options, providers=providers)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniconda/base/envs/jax2onnx2/lib/python3.12/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 472, in __init__
    self._create_inference_session(providers, provider_options, disabled_optimizers)
  File "/opt/homebrew/Caskroom/miniconda/base/envs/jax2onnx2/lib/python3.12/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 561, in _create_inference_session
    sess.initialize_session(providers, provider_options, disabled_optimizers)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Node (Gather_rev_9) Op (Gather) [ShapeInferenceError] axis must be in [-r, r-1]

This is the new reduced model:

couette.onnx.zip