[Shreya Boyane (Sonica-B)]

torch.onnx.export with its defaults (dynamo=True, optimize=True) runs the default rewrite rules, and FuseMaxMinToClip / FuseMinMaxToClip change the output shape of torch.clamp when its bounds are tensors with more dimensions than the input. A per-tensor scale such as torch.clamp(x.abs().amax(), lo, hi), with lo and hi registered as shape-[1] buffers, returns shape (1,) in PyTorch and in the unoptimized export, but a 0-d tensor after optimization, while the graph output still declares [1]. onnxscript.optimizer.optimize() does the same to any Min(Max(X, lb), ub) or Max(Min(X, ub), lb) chain with such bounds.

torchlib lowers aten::clamp.Tensor to Max followed by Min. The two rules accept any single-element bound (np.size(v) == 1) and emit Clip(X, lb, ub) with scalar bounds. Min and Max broadcast X against a float[1] or float[1, 1] bound, but Clip keeps the shape of X. When a bound has no more dimensions than X, broadcasting cannot change the shape, so only that case is wrong.

The fix is in _FuseMinMaxBase.check: the Clip fusion is skipped when a bound has more dimensions than X, or when the rank of X is unknown and the bound is not 0-d. Bounds whose rank is at most the rank of X still fuse, so the existing tests (which use float[1] bounds on 3-D and 4-D inputs) are unchanged.

flowchart LR
  subgraph main["main"]
    A["Max(X: float[], lo: float[1])"] --> B["Min(., hi: float[1])<br/>output float[1]"]
    B -. "FuseMaxMinToClip" .-> C["Clip(X, lo, hi)<br/>output float[]"]
  end
  subgraph pr["this PR"]
    D["Max(X: float[], lo: float[1])"] --> E["Min(., hi: float[1])<br/>output float[1]"]
    E -. "check fails: lo has more dimensions than X" .-> F["graph unchanged"]
  end

Loading

Repro (torch 2.15.0+cpu, onnxruntime 1.23.0):

import torch, onnxruntime as ort

class AmaxScale(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("lo", torch.tensor([1e-4]))
        self.register_buffer("hi", torch.tensor([448.0]))

    def forward(self, x):
        return torch.clamp(x.abs().amax(), self.lo, self.hi)

m, x = AmaxScale(), torch.tensor([[0.5, -3.0], [2.0, 1.0]])
print("torch:", tuple(m(x).shape))
for optimize in (False, True):
    prog = torch.onnx.export(m, (x,), dynamo=True, optimize=optimize)
    sess = ort.InferenceSession(prog.model_proto.SerializeToString())
    y = sess.run(None, {sess.get_inputs()[0].name: x.numpy()})[0]
    print(f"optimize={optimize}:", [n.op_type for n in prog.model.graph],
          "declared", prog.model.graph.outputs[0].shape, "ORT shape", y.shape)

Before (main @ 682b59d):

torch: (1,)
optimize=False: ['Abs', 'ReduceMax', 'CastLike', 'Max', 'CastLike', 'Min'] declared [1] ORT shape (1,)
optimize=True: ['Abs', 'ReduceMax', 'Clip'] declared [1] ORT shape ()

After (this branch):

torch: (1,)
optimize=False: ['Abs', 'ReduceMax', 'CastLike', 'Max', 'CastLike', 'Min'] declared [1] ORT shape (1,)
optimize=True: ['Abs', 'ReduceMax', 'Max', 'Min'] declared [1] ORT shape (1,)

On main, onnxruntime also warns when it loads the optimized model: Error merging shape info for output. 'clamp' source:{} target:{1}. Falling back to lenient merge. The values are the same in both cases (3.0). Only the rank changes, which breaks consumers that depend on it, such as Concat, Reshape or a caller comparing against the PyTorch output.

Tests: TestClipFusionBoundRank in _min_max_to_clip_test.py checks both rules with a 1-D input and float[1, 1] bounds and with a 0-d input and float[1] bounds (the condition must fail), and checks that a 1-D input with float[1] bounds still fuses and stays numerically equal. The four failure cases fail against the rule on main, because the condition passes there, and pass on this branch.

Run on Windows 11, Python 3.11, onnx 1.18.0, onnxruntime 1.23.0:

• pytest onnxscript/rewriter/rules/common -q -n 4: 355 passed, 1 skipped (the existing skip for [Broken CI] test_llama_p0_rule_set_slice_split fails #1642)
• pytest onnxscript/optimizer -q -n 4: 87 passed
• lintrunner on the two changed files: no lint issues

Note: optimizer.optimize() infers shapes before the rewrite rules run, so a graph whose intermediate values carry no value_info still fuses there. Only a bare rewriter.rewrite() on a graph without shape information now keeps Min/Max when a bound is not 0-d.