Onnx Runtime Error

[burakssen]

For this mul operation I get an error on onnxruntime:

Error: [ONNXRuntimeError] : 1 : FAIL : Load model from 02_sod_shock_tube/sod.onnx failed:Node (mul_31) Op (Mul) [ShapeInferenceError] Incompatible dimensions

also I get these warnings:

2025-07-07 22:19:58.109663 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_197' source:{201,6,210,1,1} target:{5,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.113437 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1159' source:{201,6,210,1,1} target:{5,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.117738 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2126' source:{201,6,210,1,1} target:{5,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127518 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_294' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127596 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_285' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127647 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_281' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127686 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_279' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127751 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_274' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127780 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_270' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127859 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_265' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.127913 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_256' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131508 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1256' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131566 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1247' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131606 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1243' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131635 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1241' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131680 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1236' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131705 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1232' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131764 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1227' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.131804 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1218' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135620 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2223' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135670 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2214' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135704 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2210' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135734 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2208' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135777 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2203' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135801 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2199' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135859 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2194' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.135898 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2185' source:{1,6,6,1,1} target:{1,201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146256 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_295' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146355 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_286' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146403 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_282' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146439 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_280' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146495 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_275' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146523 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_271' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146592 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_266' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.146640 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_257' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.149889 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1257' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.149945 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1248' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.149984 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1244' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.150013 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1242' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.150058 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1237' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.150083 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1233' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.150145 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1228' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.150186 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_1219' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.153850 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2224' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.153899 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2215' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.153934 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2211' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.153962 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2209' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.154003 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2204' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.154029 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2200' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.154085 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2195' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.
2025-07-07 22:19:58.154122 [W:onnxruntime:, graph.cc:117 MergeShapeInfo] Error merging shape info for output. 'var_2186' source:{6,6,1,1} target:{201,6,1,1}. Falling back to lenient merge.

I am not really sure what's causing this, I am using the same code with different inputs and I get this not really sure why.

If you want, I can try to create a simple example but not really sure how to cause the same error with a minimal setup.

[enpasos]

The error you're hitting on the Mul node (mul_31) is likely due to incompatible shapes that aren't broadcastable under ONNX's stricter rules, even though they look compatible visually (201×6×1×1). One of the inputs might still carry symbolic or mismatched dimensions introduced earlier (e.g., via Slice, Squeeze, Expand, or Reshape), and ONNX Runtime shape inference doesn’t always handle these leniently.

The many [W: graph.cc:117 MergeShapeInfo] warnings you're seeing confirm that shape inconsistencies are present — ONNX Runtime is attempting to fall back to lenient merging, but it’s a sign something upstream isn’t fully aligning.

To better understand and fix the issue, it would be super helpful if you could share a way to reproduce this — even if it’s not minimal or simplified. A full working code snippet that leads to this model being generated and saved would let me directly debug the issue in context and ensure that any fix fully covers your case.

Let me know if you’re able to share that.

[burakssen]

import os
import numpy as np
import jax
import jax.numpy as jnp
from jax2onnx import to_onnx

from jaxfluids import InputManager, InitializationManager, SimulationManager
from jaxfluids.feed_forward.feed_forward import (
    FeedForwardSetup, 
    MachineLearningSetup, 
    CallablesSetup, 
    ParametersSetup
)

# Configuration
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
jax.config.update("jax_enable_x64", True)

# Constants
INITIAL_TIME = 0.0
TIME_STEP = 1e-3
OUTPUT_FILE = "sod.onnx"

def setup_simulation():
    """Initialize simulation components and return necessary buffers."""
    input_manager = InputManager("sod.json", "numerical_setup.json")
    initialization_manager = InitializationManager(input_manager)
    sim_manager = SimulationManager(input_manager)
    
    jxf_buffers = initialization_manager.initialization()    
    # Extract initial primitives
    primitives = jxf_buffers.simulation_buffers.material_fields.primitives
    nhx, nhy, nhz = sim_manager.domain_information.domain_slices_conservatives
    primes_init = primitives[..., nhx, nhy, nhz]
    
    return sim_manager, primes_init

def create_feed_forward_fn(sim_manager):
    """Create the feed forward function for ONNX conversion."""
    def feed_forward_fn(y_current, t_current, time_step):
        time = t_current[0]
        time_step = time_step[0]
        
        solution_array, times_array = sim_manager._feed_forward(
            y_current,
            time_step,
            time,
            FeedForwardSetup(1, 100, True, False, False, False, False)
        )
        
        y_current = solution_array["primitives"][-1]
        time = times_array
        
        return y_current, time
    
    return feed_forward_fn

def convert_to_onnx(feed_forward_fn, primes_init):
    """Convert the feed forward function to ONNX format."""
    inputs = [
        primes_init,
        jnp.array([INITIAL_TIME], dtype=jnp.float64),
        jnp.array([TIME_STEP], dtype=jnp.float64)
    ]
    
    model = to_onnx(
        feed_forward_fn,
        inputs=inputs,
        enable_double_precision=True
    )
    
    with open(OUTPUT_FILE, "wb") as f:
        f.write(model.SerializeToString())
    
    print(f"Model saved to {OUTPUT_FILE}")

def main():
    """Main execution function."""
    sim_manager, primes_init = setup_simulation()
    feed_forward_fn = create_feed_forward_fn(sim_manager)
    convert_to_onnx(feed_forward_fn, primes_init)

if __name__ == "__main__":
    main()

Input files:

numerical_setup.json
sod.json

Ok so this is the code that I am trying to generate onnx model from. We are using our own framework for this. You can find it on https://github.com/tumaer/JAXFLUIDS

I was not really sure about what to do really, if there is anything comes to your mind that we need to fix on our end just let me know, and thank you for your help in advance.

[burakssen]

@enpasos I wanted to let you know for now we found a work around for this. We had another implementations for the same thing with different input, I guess in this implementation it didn't work as expected.

[enpasos]

@burakssen With v0.7.1, we’ve fixed a bug in the test system that now ensures numeric equivalence between ONNX-generated functions and their JAX counterparts. I’m taking a week off now but plan to look into your issue once I’m back.

[enpasos]

Update: branch onnx_ir_and_jax_072 now integrates JAX 0.7.2, Flax 0.12.0, Equinox 0.13.1, and onnx-ir 0.1.10 (rough cut, WIP). I’ve started tackling the issues on this base: first step was serializing the feed-forward closed jaxpr into a portable payload using export_issue52_payload.py. A standalone repro is in jax2onnx/sandbox/issue52_scatter_payload_repro.py + issue52_feedforward_payload.npz. From there, I’m building test cases and iterating — some progress, but no fix yet (see work_notes_issue52.md).