Merge output shape with input shape instead of override

[wodesuck]

_constant_folding.cast override output.shape with input.shape, that may make a static shape to dynamic shape. Here should use _merge_shapes instead.

[Copilot]

Pull Request Overview

This PR fixes a shape handling issue in the constant folding optimization for Cast operations. The change prevents static shapes from being incorrectly converted to dynamic shapes by merging output and input shapes instead of overriding them.

• Replaces direct shape override with intelligent shape merging logic
• Preserves static shape information when possible during Cast operations

[justinchuby]

@gramalingam @titaiwangms

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.04%. Comparing base (dddf0c2) to head (4b79f9b).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2578   +/-   ##
=======================================
  Coverage   70.04%   70.04%           
=======================================
  Files         223      223           
  Lines       26215    26213    -2     
  Branches     2583     2582    -1     
=======================================
- Hits        18363    18362    -1     
  Misses       6946     6946           
+ Partials      906      905    -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[justinchuby]

Could you share a concrete example when you see this case?

[wodesuck]

import torch
import torch.nn as nn


class Model(nn.Module):
    def forward(self, x):
        return x.new_zeros(x.shape, dtype=torch.int64)


def main():
    model = Model()
    args = torch.rand(4, 4),
    batch = torch.export.Dim("batch")
    dynamic_shapes = {"x": {0: batch}}
    torch.onnx.export(
        model,
        args,
        "model_test.onnx",
        dynamic_shapes=dynamic_shapes,
        dynamo=True,
    )


if __name__ == "__main__":
    main()

The output shape should be batch x 4, but now it's ? x ?.

[justinchuby]

Thanks!