Feature Request: Support setting inputs and outputs as NCHW to prevent having to transpose data several times through the model

[Artoriuz]

The preferred data format for images on JAX is NHWC, but ONNX uses NCHW instead. When exporting a model directly, this results in many operations getting "wrapped around" by extra operations transposing the data, like this:

This was also the case when exporting from Tensorflow, which also defaults to NHWC, but the TF/Keras exporter (https://github.com/onnx/tensorflow-onnx/blob/4fed7de9534b6a084f7f2326bae775545bd97f9e/tf2onnx/convert.py#L412) supports the following arguments:

inputs_as_nchw: transpose inputs in list from nhwc to nchw
outputs_as_nchw: transpose outputs in list from nhwc to nchw

With these enabled, the graph looked clean and mostly free of the unnecessary transpose operations as everything could be treated as NCHW instead.

It would be nice if jax2onnx had a similar mechanism.

Thanks in advance.

[enpasos]

Thanks for the detailed report and the screenshot!

To ensure we implement exactly what you need, I’d like to separate this into two distinct categories, as jax2onnx handles them differently:

1. Graph Internals (The Middle): We already have an optimization pass intended to condense redundant operations (like back-and-forth transposes) inside the model. If you are seeing unnecessary transposes deep within the graph structure, that would be a bug or a gap in our current optimization logic that we should fix directly.
2. Graph Boundaries (Inputs/Outputs): The transposes immediately following inputs and preceding outputs are currently intentional. Since JAX natively uses NHWC (Channels Last) and ONNX operators generally expect NCHW, we must wrap the graph to ensure the ONNX model respects the exact I/O contract of the JAX function.

The Feature Plan:
We agree that for deployment (e.g., TensorRT), you often want to break that strict I/O contract to get a clean NCHW interface. We can support this by adding a "boundary rewrite" feature (e.g., inputs_as_nchw, outputs_as_nchw flags) to to_onnx.

Request for Reproduction:
However, before we implement that feature, could you please provide a minimal code example or a small reproduction script that generates the graph you are seeing?

This will help us confirm whether your "messy graph" is purely due to these boundary contracts (which the new feature will solve) or if we are also missing internal optimizations (which we should fix as bugs).