def simple_foldable(): # Graph: # c = (a + b) # output = input + c # Should fold to: # output = input + c weights = np.ones(shape=(1, 3), dtype=np.float32) graph = Graph() inp = Variable("input", shape=(1, 3), dtype=np.float32) c = graph.add(weights, weights, name="c") out = graph.add(inp, c) graph.inputs = [inp] graph.outputs = [out] yield graph
def test_shape_gather(self, shape, indices): indices = np.array(indices) inp = Variable("input", dtype=np.float32, shape=shape) graph = Graph(inputs=[inp]) inp_shape = graph.shape(inp) shape_part = graph.gather(inp_shape, indices=indices) graph.outputs = [ graph.add(shape_part, shape_part), graph.gather(inp_shape, indices=[0]), graph.gather(inp_shape, indices=np.array(0)), ] graph.fold_constants() if shape is not None: assert isinstance(graph.outputs[0], Constant) expected_shape = np.array(shape)[indices].astype(np.int64) * 2 assert np.all(graph.outputs[0].values == expected_shape) else: assert isinstance(graph.outputs[0], Variable) assert isinstance(graph.outputs[1], Variable) assert isinstance(graph.outputs[2], Variable)
def one_hop_foldable(): # Graph: # c = (a + b) # e = (c + d) # output = input + e # Should fold to: # output = input + e weights = np.ones(shape=(1, 3), dtype=np.float32) graph = Graph() inp = Variable("input", shape=(1, 3), dtype=np.float32) c = graph.add(weights, weights, name="c") e = graph.add(c, weights, name="e") out = graph.add(inp, e) graph.inputs = [inp] graph.outputs = [out] yield graph
def test_register(self): @Graph.register() def add(self, a, b): return self.layer(op="Add", inputs=[a, b], outputs=["add_out"]) graph = Graph() [output] = graph.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 add(self, a, b): return self.layer(op="Add", inputs=[a, b], outputs=["add_out"]) @Graph.register(opsets=[10]) def add(self, a, b): return self.layer(op="Add-10", inputs=[a, b], outputs=["add_out"]) graph = Graph() [output] = graph.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.add("a", "b") assert "add_out" in output.name assert len(graph_opset10.nodes) == 1 assert graph_opset10.nodes[-1].op == "Add-10"
def foldable_with_invalid_node(): # Graph # c = (a + b) # e = fake(d) # f = (e + c) # out = inp + f # # c should be folded even though e is the output of an # invalid node. weights = np.ones(shape=(1, 3), dtype=np.float32) graph = Graph() inp = Variable("input", shape=(1, 3), dtype=np.float32) c = graph.add(weights, weights, name="c") e = graph.fake(weights, name="e") f = graph.add(e, c, name="f") out = graph.add(inp, f, name="output") graph.inputs = [inp] graph.outputs = [out] yield graph
def test_const_inp_but_non_foldable_nested_graph(self): cond = gs.Constant("cond", values=np.array(True)) X = gs.Variable("X", dtype=np.float32, shape=(1, )) graph = Graph(inputs=[X]) then_graph = Graph(name="Then") then_graph.outputs = [then_graph.add(X, X)] else_graph = Graph(name="Else") else_graph.outputs = [else_graph.add(X, else_graph.add(X, X))] # Even though if_op looks foldable because it has all constant inputs, # it's not, since its subgraphs depend on variables in the outer scope. graph.outputs = [graph.if_op(cond, then_graph, else_graph)] # This should not raise because the `If` node should be excluded from # constant folding. graph.fold_constants(error_ok=False).cleanup() assert graph.nodes[0].op == "If" assert len(then_graph.nodes) == 1 assert len(else_graph.nodes) == 2
def test_with_nested_graph(self): cond = gs.Variable("cond", dtype=np.bool, shape=(1, )) X = gs.Variable("X", dtype=np.float32, shape=(1, )) Y = gs.Constant("Y", values=np.ones((1, ), dtype=np.float32)) graph = Graph(inputs=[X, cond]) then_graph = Graph(name="Then") then_graph.outputs = [then_graph.add(Y, Y)] else_graph = Graph(name="Else") else_graph.outputs = [else_graph.add(X, else_graph.add(Y, Y))] graph.outputs = [graph.if_op(cond, then_graph, else_graph)] graph.fold_constants() graph.cleanup() assert len(then_graph.nodes) == 0 assert np.all(then_graph.outputs[0].values == (Y.values * 2)) assert len(else_graph.nodes) == 1 assert isinstance(else_graph.nodes[0].inputs[1], Constant) assert np.all(else_graph.nodes[0].inputs[1].values == (Y.values * 2))
def test_input_is_output(self): graph = Graph() A = Variable("A", dtype=np.float32, shape=(1, 1)) B = Variable("B", dtype=np.float32, shape=(1, 1)) C = graph.add(A, B) graph.inputs = [A, B] graph.outputs = [C, B, A] # Out of order w/ respect to Add node inputs # Graph should remain unchanged after cleanup, including I/O tensors. graph.cleanup() assert graph.inputs == [A, B] assert graph.outputs == [C, B, A] assert len(graph.nodes) == 1 assert graph.nodes[0].inputs == [A, B] assert graph.nodes[0].outputs == [C]
def test_node_used_only_in_nested_graph(self): X = Variable("X", dtype=np.float32, shape=(1, )) Y = Variable("Y", dtype=np.float32, shape=(1, )) graph = Graph(inputs=[X, Y]) X_p = graph.identity( X) # X_p is only used by the subgraph, not in the outer graph. subgraph_inp = Variable("subgraph_input", dtype=np.float32, shape=(1, )) subgraph = Graph(inputs=[subgraph_inp]) subgraph.outputs = [subgraph.add(subgraph_inp, X_p)] graph.outputs = [graph.nested(Y, subgraph)] graph.cleanup(remove_unused_graph_inputs=True) assert graph.nodes[0].op == "Identity" assert graph.nodes[0].inputs == [X]