def test_layer_with_attrs(self): graph = Graph() outputs = graph.layer(op="Add", name="node", attrs={"fake_attr": 0}) assert len(graph.nodes) == 1 assert graph.nodes[-1].op == "Add" assert graph.nodes[-1].name == "node" assert graph.nodes[-1].attrs["fake_attr"] == 0
def test_layer_with_strings(self): x0 = "x0" x1 = "x1" y0 = "y0" y1 = "y1" graph = Graph() outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1]) assert len(graph.nodes) == 1 assert [prefix in tensor.name for prefix, tensor in zip([x0, x1], graph.nodes[-1].inputs)] assert [prefix in tensor.name for prefix, tensor in zip([y0, y1], graph.nodes[-1].outputs)] assert graph.nodes[-1].outputs == outputs
def test_layer_with_tensors(self): x0 = Variable("x0") x1 = Variable("x1") y0 = Variable("y0") y1 = Variable("y1") graph = Graph() outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1]) assert outputs == [y0, y1] assert len(graph.nodes) == 1 assert graph.nodes[-1].inputs == [x0, x1] assert graph.nodes[-1].outputs == outputs
def test_layer_with_arrays(self): x0 = np.array([1]) x1 = np.array([1]) y0 = "y0" y1 = "y1" graph = Graph() outputs = graph.layer(op="Fake", inputs=[x0, x1], outputs=[y0, y1]) assert [prefix in tensor.name for prefix, tensor in zip([y0, y1], graph.nodes[-1].outputs)] assert len(graph.nodes) == 1 assert graph.nodes[-1].inputs[0].values == x0 assert graph.nodes[-1].inputs[1].values == x1 assert graph.nodes[-1].outputs == outputs