ONNX export of CumSum produces different data type

[mergian]

🐛 Describe the bug

When running a model with torch.cumsum(..., dtype=None) where dtype is set to None, an integer data type gets promoted to torch.int64. However, if the same model gets exported to ONNX, the dtype does not get promoted. See following example:

import torch
import onnx

class Model(torch.nn.Module):
    def forward(self, x):
        return x.cumsum(0, dtype=None), x.cumsum(0, dtype=torch.int32), x.cumsum(0, dtype=torch.float)

input = torch.rand(1, 2, 3).to(torch.int32)
model = Model()
output = model(input)

torch.onnx.export(model, input, 'error.onnx', opset_version=11)

omodel = onnx.load('error.onnx')
lookup = {
    onnx.TensorProto.BOOL:      'onnx.TensorProto.BOOL',
    onnx.TensorProto.DOUBLE:    'onnx.TensorProto.DOUBLE',
    onnx.TensorProto.FLOAT16:   'onnx.TensorProto.FLOAT16',
    onnx.TensorProto.FLOAT:     'onnx.TensorProto.FLOAT',
    onnx.TensorProto.INT8:      'onnx.TensorProto.INT8',
    onnx.TensorProto.INT16:     'onnx.TensorProto.INT16',
    onnx.TensorProto.INT32:     'onnx.TensorProto.INT32',
    onnx.TensorProto.INT64:     'onnx.TensorProto.INT64',
    onnx.TensorProto.UINT8:     'onnx.TensorProto.UINT8',
    onnx.TensorProto.UINT16:    'onnx.TensorProto.UINT16',
    onnx.TensorProto.UINT32:    'onnx.TensorProto.UINT32',
    onnx.TensorProto.UINT64:    'onnx.TensorProto.UINT64'
}

print('PyTorch Output DTypes: {}'.format(tuple(o.dtype for o in output)))
print('ONNX Output DTypes: {}'.format(
    tuple(lookup.get(o.type.tensor_type.elem_type) for o in omodel.graph.output))
)

Output is:

PyTorch Output DTypes: (torch.int64, torch.int32, torch.float32)
ONNX Output DTypes: ('onnx.TensorProto.INT32', 'onnx.TensorProto.INT32', 'onnx.TensorProto.FLOAT')

As you can see, in the dtype=None case, PyTorch uses INT64 while ONNX uses INT32.

Versions

Collecting environment information...
PyTorch version: 1.11.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 10.3.0
Clang version: Could not collect
CMake version: version 3.23.1
Libc version: glibc-2.17

Python version: 3.7.12 (default, Feb 6 2022, 20:29:18) [GCC 10.2.1 20210130 (Red Hat 10.2.1-11)] (64-bit runtime)
Python platform: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-centos-7.9.2009-Core
Is CUDA available: False
CUDA runtime version: 11.3.109
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] torch==1.11.0
[pip3] torchvision==0.12.0
[conda] Could not collect

[justinchuby]

Thanks for reporting this issue! We have added it to our list for investigation.