Add float16 support for cumsum

If I use cumsum in a model with float16 types, I get this error when I try to import the model:

onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model. Type Error: Type 'tensor(float16)' of input parameter (input) of operator (CumSum) in node (CumSum_1) is invalid.

If I use topk in a model with float16 types, I get this error when I try to run inference on the model:

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running TopK node. Name:'TopK_3' Status Message: Type not supported for TopK operator

Here is a minimal chunk of code designed to construct a model, export the model, import the model, and run inference on the model:

import torch
from torch import tensor
import onnxruntime as ort
import os

class TopK(torch.nn.Module):
    def __init__(self, k):
        super().__init__()
        self.k = k

    def forward(self, x):
        results = torch.topk(x, self.k)
        return results.indices

class CumSum(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.cumsum(x, 0)

# change the dtype to test another dtype
x = torch.randn(10, dtype=torch.float16, device='cuda')

# uncomment this line to test topk
model = TopK(5)

# uncomment this line to test cumsum
# model = CumSum()

print('torch inference:')
print(model(x))

print('exporting...')
os.makedirs('simple', exist_ok=True)
fn = 'simple/model.onnx'
with torch.no_grad():
    torch.onnx.export(
        model, x,
        f=fn,
        opset_version=12,
        input_names=['input'],
        output_names=['output']
    )

print('importing...')
options = ort.SessionOptions()
session = ort.InferenceSession(fn, options, providers=['CUDAExecutionProvider'])

print('onnx inference:')
result = session.run(['output'], {'input': x.cpu().numpy()})
print(result)

For examples of the expected results, here is testing of topk with float32:

torch inference:
tensor([4, 9, 5, 1, 8], device='cuda:0')
exporting...
importing...
onnx inference:
[array([4, 9, 5, 1, 8], dtype=int64)]

And here is testing cumsum with float32:

torch inference:
tensor([-0.3108, -0.4656,  0.0363,  1.1375,  0.8920,  0.9415, -0.2531,  2.2555,
         2.7834,  2.3672], device='cuda:0')
exporting...
importing...
onnx inference:
[array([-0.3108134 , -0.4655881 ,  0.03631312,  1.1374552 ,  0.89203495,
        0.9414741 , -0.25308657,  2.2555485 ,  2.7834275 ,  2.3671858 ],
      dtype=float32)]

I am using CUDA 10.2 and:

torch             1.5.0
onnx              1.7.0
onnxruntime-gpu   1.3.0
onnxruntime-tools 1.3.0.1007

I can report this as two different issues if that helps.

[linkerzhang]

as the error said, the “float16” type has not been supported in onnx runtime, though it’s supported in the op definition (both comsum and topk).

You may report this issue in onnx runtime community (http://github.com/microsoft/onnxruntime) to get more help/support. Adding @snnn @faxu .

[jcwchen]

Cumsum does support float16 now: https://github.com/onnx/onnx/pull/3195. Please reopen it if you have other questions.