Skip to content

Commit 5fbb6e7

Browse files
Merge pull request #1274 from abhishek-gola:size_layer_add
* Added models for test_size 0d and 1d * added model generation script
1 parent 5907af6 commit 5fbb6e7

File tree

3 files changed

+21
-0
lines changed

3 files changed

+21
-0
lines changed
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import onnx
2+
from onnx import helper, TensorProto
3+
4+
def make_size_model(model_path, name, input_shape):
5+
# input_shape: e.g. ['N'] for 1D dynamic, or [] for scalar
6+
inp = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape if input_shape else [])
7+
out = helper.make_tensor_value_info("y", TensorProto.INT64, [])
8+
node = helper.make_node("Size", ["x"], ["y"])
9+
graph = helper.make_graph([node], name, [inp], [out])
10+
opset = helper.make_operatorsetid("", 13)
11+
model = helper.make_model(graph, opset_imports=[opset], producer_name="size-tests")
12+
onnx.checker.check_model(model)
13+
onnx.save(model, model_path)
14+
15+
make_size_model("test_size_1d_model.onnx",
16+
"test_size_1d", ["N"])
17+
18+
# True scalar input (rank-0, shape []).
19+
make_size_model("test_size_0d_model.onnx",
20+
"test_size_0d", [])
21+
print("Wrote models.")
76 Bytes
Binary file not shown.
81 Bytes
Binary file not shown.

0 commit comments

Comments
 (0)