def test_replace(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) with pytest.raises(Exception, match="Cannot replace.*"): var4.fgraph = object() # Trigger a `FunctionGraph` ownership error fg.replace(var4, var1, verbose=True) var4.fgraph = fg with pytest.raises(BadOptimization): var0 = MyVariable2("var0") # The types don't match and one cannot be converted to the other fg.replace(var3, var0) # Test a basic replacement fg.replace_all([(var3, var1)]) assert var3 not in fg.variables assert fg.apply_nodes == {var4.owner, var5.owner} assert var4.owner.inputs == [var1, var2]
def test_replace_circular(self): """`FunctionGraph` allows cycles--for better or worse.""" 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) fg.replace_all([(var3, var4)]) # The following works (and is kind of gross), because `var4` has been # mutated in-place assert fg.apply_nodes == {var4.owner, var5.owner} assert var4.owner.inputs == [var4, var2]