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"
