[ONNX] add cast operator after reduce to match desired dtype

[TaiPhamD]

This PR conditionally inserts a cast operator after a reduction operation to match the specified dtype in the exported ONNX model. The code changes affect opset9, and opset13.

I understand there's an automatic upcast to int64 before reduction most likely to prevent overflow so I left that alone and only conditionally add casting back to desired dtype.

Test int32

import torch
import onnx
a = torch.tensor([10, 20, 30, 80], dtype=torch.int32)
def test():
    class SumInt32(torch.nn.Module):
        def forward(self, a):
            return torch.sum(a, dtype=torch.int32)

    sumi = SumInt32().eval()
    assert sumi(a).dtype == torch.int32
    print("Torch model output type matches input type")

    torch.onnx.export(sumi, (a), "/tmp/sumi_int32.onnx", opset_version=12)
    model = onnx.load("/tmp/sumi_int32.onnx")

    assert model.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.INT32
    print("ONNX model output type matches input type")
test()

Test int64

import onnx
import torch

a = torch.tensor([10, 20, 30, 80], dtype=torch.int64)


def test():
    class SumInt64(torch.nn.Module):
        def forward(self, a):
            return torch.sum(a, dtype=torch.int64)

    sumi = SumInt64().eval()
    assert sumi(a).dtype == torch.int64
    print("Torch model output type matches input type")
    torch.onnx.export(sumi, (a), "/tmp/sumi_int64.onnx", opset_version=12)
    model = onnx.load("/tmp/sumi_int64.onnx")
    assert model.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.INT64
    print("ONNX model output type matches input type")


test()

Fixes #100097

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100700

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 1 Pending

As of commit 274c2a8:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: TaiPhamD (42696e5)

[TaiPhamD]

I found a few issues with this one so will close until i fix it then will re-open.

[thiagocrepaldi]

Thank you for the PR. I have pointed out an issue in which how you accessed the scalar tpe for a torch._C.Value. Please use the recommended way

[thiagocrepaldi]

parsing type from name is the #1 source of problems (because of the previous comment). Please only use from_value or from_dtype.

[thiagocrepaldi]

This can cause segfault/assert crashes when model is torch.jit.script.
To prevent this, use exclusively JitScalarType's public APIs to extract type information from torch._C.Value nodes in a safe way.

[TaiPhamD]

thank you I will make that change.

[TaiPhamD]

sorry for the lint issue. I used flake8 to test for lint issues locally but just realized the CI uses lintrunner so I'll make that change.

[thiagocrepaldi]

sorry for the lint issue. I used flake8 to test for lint issues locally but just realized the CI uses lintrunner so I'll make that change.

I always do make setup_lint, lintrunner init and lintrunner -a before pushing a PR to make sure linting is fine

[BowenBao]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here