Spaces:
Sleeping
Sleeping
File size: 1,263 Bytes
dc2106c |
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 |
# Copyright (c) ONNX Project Contributors
# SPDX-License-Identifier: Apache-2.0
import unittest
import onnx
from onnx import parser, printer
class TestBasicFunctions(unittest.TestCase):
def check_graph(self, graph: onnx.GraphProto) -> None:
self.assertEqual(len(graph.node), 3)
self.assertEqual(graph.node[0].op_type, "MatMul")
self.assertEqual(graph.node[1].op_type, "Add")
self.assertEqual(graph.node[2].op_type, "Softmax")
def test_parse_graph(self) -> None:
text0 = """
agraph (float[N, 128] X, float[128,10] W, float[10] B) => (float[N] C)
{
T = MatMul(X, W)
S = Add(T, B)
C = Softmax(S)
}
"""
graph1 = parser.parse_graph(text0)
text1 = printer.to_text(graph1)
graph2 = parser.parse_graph(text1)
text2 = printer.to_text(graph2)
# Note that text0 and text1 should be semantically-equivalent, but may differ
# in white-space and other syntactic sugar. However, we expect text1 and text2
# to be identical.
self.assertEqual(text1, text2)
self.check_graph(graph2)
if __name__ == "__main__":
unittest.main()
|