def test_register(self): @Graph.register() def fake_add(self, a, b): return self.layer(op="Add", inputs=[a, b], outputs=["add_out"]) graph = Graph() [output] = graph.fake_add("a", "b") assert "add_out" in output.name assert len(graph.nodes) == 1 assert graph.nodes[-1].op == "Add"
def test_register_opset(self): @Graph.register(opsets=[11]) def fake_add(self, a, b): return self.layer(op="Add", inputs=[a, b], outputs=["add_out"]) @Graph.register(opsets=[10]) def fake_add(self, a, b): return self.layer(op="Add-10", inputs=[a, b], outputs=["add_out"]) graph = Graph() [output] = graph.fake_add("a", "b") assert "add_out" in output.name assert len(graph.nodes) == 1 assert graph.nodes[-1].op == "Add" graph_opset10 = Graph(opset=10) [output] = graph_opset10.fake_add("a", "b") assert "add_out" in output.name assert len(graph_opset10.nodes) == 1 assert graph_opset10.nodes[-1].op == "Add-10"