def test_init(self): var1 = MyVariable("var1") var2 = MyVariable("var2") var3 = op1(var1) var4 = op2(var3, var2) fg = FunctionGraph([var1, var2], [var3, var4], clone=False) assert fg.inputs == [var1, var2] assert fg.outputs == [var3, var4] assert fg.apply_nodes == {var3.owner, var4.owner} assert fg.update_mapping is None assert fg.check_integrity() is None assert fg.variables == {var1, var2, var3, var4} assert fg.clients(var1) == [(var3.owner, 0)] assert fg.clients(var2) == [(var4.owner, 1)] assert fg.clients(var3) == [(var4.owner, 0), ("output", 0)] assert fg.clients(var4) == [("output", 1)]
def test_remove_client(self): var1 = MyVariable("var1") var2 = MyVariable("var2") var3 = op1(var2, var1) var4 = op2(var3, var2) var5 = op3(var4, var2, var2) fg = FunctionGraph([var1, var2], [var3, var5], clone=False) assert fg.variables == {var1, var2, var3, var4, var5} assert fg.clients(var2) == [ (var3.owner, 0), (var4.owner, 1), (var5.owner, 1), (var5.owner, 2), ] fg.remove_client(var2, (var4.owner, 1)) assert fg.clients(var2) == [ (var3.owner, 0), (var5.owner, 1), (var5.owner, 2), ] fg.remove_client(var1, (var3.owner, 1)) assert fg.clients(var1) == [] assert var4.owner in fg.apply_nodes # This next `remove_client` should trigger a complete removal of `var4`'s # variables and `Apply` node from the `FunctionGraph`. # # Also, notice that we already removed `var4` from `var2`'s client list # above, so, when we completely remove `var4`, `fg.remove_client` will # attempt to remove `(var4.owner, 1)` from `var2`'s client list again. # This attempt would previously raise a `ValueError` exception, because # the entry was not in the list. fg.remove_client(var4, (var5.owner, 0), reason="testing") assert var4.owner not in fg.apply_nodes assert var4.owner.tag.removed_by == ["testing"] assert not any(o in fg.variables for o in var4.owner.outputs)