A module ran correctly in PyTorch, but export stopped at a Python round() call that depended on the input dimensions. The expression was straightforward: calculate a pooling output size from a feature map's aspect ratio, then pass the integer dimensions to the next operator.
After fixing the input size, moving that calculation outside the graph allowed export and comparison with ONNX Runtime. This case concerns how Python shape calculations enter a captured graph, and which input conditions a static deployment must preserve.
1. Running in Python and capturing a graph are separate steps
In ordinary eager execution, dimensions often behave as Python integers. Code can compare them, divide them, round them and choose how to call tensor operators. This general size calculation is one example:
def output_size(height, width, short_side):
if width >= height:
return short_side, round(width * short_side / height)
return round(height * short_side / width), short_side
Ordinary Python can execute it without difficulty. During export, height and width may participate as symbolic dimensions or tracing objects. Python branches, floating-point arithmetic and rounding then have to be captured, constrained or specialized. The details depend on the exporter and its version.
The blocking point in this record was that Python shape expression. It does not establish that ONNX cannot represent rounding. Python's built-in round(), a rounding operation on a tensor and an operator in an exported graph belong to different layers. Replacing the function name may still leave the downstream operator's size-argument requirements unsatisfied.
The newer torch.export path handles Python control flow and records shape constraints. Traditional tracing may instead fix decisions made for the example input. Recording options such as dynamo, dynamic_shapes and dynamic_axes helps explain what the exported file actually supports.
In the public torch.export examples, dimensions not declared dynamic are generally specialized to the example input. Dynamic dimensions require explicit declarations and must satisfy the constraints produced during capture. There are therefore two approaches to this expression: preserve its dynamic behavior and let the exporter capture and validate the constraints, or explicitly fix the supported input range and calculate the dimensions in advance. With the second approach, the application's input contract becomes part of correctness.
2. Establish whether deployment needs dynamic dimensions
If the application always supplies one resolution and the corresponding intermediate feature shapes are fixed, the pooling target can be calculated once when constructing an adapter. Its forward method then uses a known integer tuple instead of deriving it from the input through Python.
This adaptation preserves the original calculation under fixed input conditions. It must retain the original ratios, offsets and rounding rules, and reject inputs outside the supported dimensions. The following example only illustrates that fixed-shape contract:
import torch
class StaticPool(torch.nn.Module):
def __init__(self, height, width, short_side):
super().__init__()
self.expected_hw = (height, width)
self.target_hw = output_size(height, width, short_side)
def forward(self, x):
if tuple(x.shape[-2:]) != self.expected_hw:
raise ValueError("unexpected spatial shape")
return torch.nn.functional.adaptive_avg_pool2d(x, self.target_hw)
The Python check applies before export or during eager execution. An exporter may specialize or eliminate it; do not assume that it becomes a runtime check inside ONNX. The application must still validate inputs against the model's input signature.
Static specialization also does not guarantee that every pooling-size combination can be exported. Once the shape calculation is removed, pooling, interpolation and other operators still need validation with the selected exporter and target runtime.
The benefit comes from removing decisions that are already known. With fixed height and width, the long-side branch, rounded result and pooling target can all be determined in advance. Graph capture no longer has to express those Python operations. Claiming arbitrary input dimensions would conflict with the fixed values. The size check constrains the conditions under which the adaptation remains valid.
If a product needs only a few known resolutions, separate static artifacts can be generated and selected by input size. This reduces dynamic control within each graph but increases the number of artifacts and the application's routing responsibilities. A single dynamic graph requires broader shape constraints and backend coverage. The actual set of supported inputs should determine the choice; merely supplying dynamic-axis arguments does not settle it.
3. Check the static expression against the original computation
Rounding boundaries are easy to overlook. If the original expression adds an offset or treats the long and short sides differently, precomputation must preserve that behavior. Python round() handles half-integers differently from some other rounding schemes. Replacing it with truncation or another rounding operation can change the output dimensions.
For example, Python round(2.5) returns 2, whereas floor(2.5 + 0.5) returns 3. A difference of one changes the number of spatial elements after pooling and the input to a later reshape or interpolation. A small offset in the original expression can decide the result at a rounding boundary, so its small magnitude is no reason to omit it. The integer result and the surrounding shape calculations need to be checked together.
Operator semantics need a separate check. ONNX AveragePool describes sampling windows through attributes such as kernel, stride and padding. PyTorch adaptive pooling derives each output location's coverage from the target output size. An ordinary pooling configuration is equivalent only when the dimensions and windows meet the required conditions. Equal output shapes do not establish equal coverage, edge handling or divisors. Resolving shape capture does not justify replacing arbitrary adaptive pooling with an unchecked ordinary pool.
Also check subsequent reshapes, interpolation targets and layouts. Matching shapes do not prove that the numerical path is unchanged, particularly when replacing dynamic dimensions with constants also specializes nearby branches.
In this isolated adaptation, the original and static PyTorch implementations were compared at a fixed feature shape. The measured maximum absolute differences for FP64 forward outputs, input gradients and shared-parameter gradients were all zero. The subsequent ONNX Runtime comparison had a relative L2 error of approximately 1.70e-7. The first set checks the rewrite under those input conditions; the second checks the exported runtime output.
The gradient checks compared the two PyTorch implementations. They did not validate ONNX training or backward execution. Zero differences describe only the tested state and inputs. Changes to resolution, ratio or other shape-dependent settings require regenerating and rechecking the adaptation.
These checks have different responsibilities. Structural validation checks graph format and operator constraints. The original-framework comparison checks whether the adapted expression preserves the tested computation. The ORT comparison checks exported execution. Passing one does not replace the others. A legal graph with an incorrect size constant may pass a checker and still fail a numerical comparison.
4. Confirm the supported scope after export
A fixed-spatial-shape ONNX graph can retain Shape nodes because it may still read the batch dimension or other dimensions. Their presence does not demonstrate dynamic input support. Conversely, naming dynamic axes does not prove that every shape-dependent branch is dynamic.
Inspect the input signature, graph constants and shape constraints alongside actual execution. If dynamic height and width are required, test rectangular inputs, swapped long and short sides, rounding boundaries and unsupported sizes. Branches and shape calculations must use expressions the exporter can capture. One passing square input does not cover these behaviors.
Shape inference also has a limited role. It propagates types and dimensions that can be determined from the graph; it is not a numerical proof for every possible input. Two graphs can infer equal output dimensions while their pooling or interpolation differs. Some dimension relationships hold only after concrete inputs or extra constraints are supplied.
A static artifact should therefore document its input signature, construction settings and preprocessing assumptions. Checking only the outer image dimensions is insufficient when cropping, padding or other preprocessing affects intermediate feature shapes. Caching and export share a requirement: precomputed values must remain traceable to input conditions that still hold.
The validated scope here was a fixed-spatial-shape expression and an ORT CPU output comparison. It resolved export at that shape without extending support to arbitrary resolutions. For intermediate-output and application-interface comparisons, see Intermediate features and numerical alignment in ONNX conversion.