1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
|
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from onnx import TensorProto
from onnx.helper import make_tensor
from onnxscript.main import script
from onnxscript.onnx_opset import opset14, opset16
from onnxscript.onnx_types import INT64
@script()
def shape_A(data, start: INT64[1], end: INT64[1]):
shape = opset16.Shape(data)
zero = opset16.Constant(value=make_tensor("zero", TensorProto.INT64, [1], [0]))
if end == zero:
length = opset16.Shape(opset16.Shape(data))
res = opset16.Slice(shape, start, length, zero)
else:
res = opset16.Slice(shape, start, end, zero)
return res
@script()
def shape_B(data, start: INT64[1], end: INT64[1]):
shape = opset14.Shape(data)
zero = opset14.Constant(value=make_tensor("zero", TensorProto.INT64, [1], [0]))
if end == zero:
length = opset14.Shape(opset14.Shape(data))
res = opset14.Slice(shape, start, length, zero)
else:
res = opset14.Slice(shape, start, end, zero)
return res
@script(default_opset=opset16)
def inc_any(data):
# The converter cannot know which opset to use unless it is specified
# in the decorator.
return data + 1
|
