def test_remove_in_and_out(self):
var1 = MyVariable("var1")
var2 = MyVariable("var2")
op1_out = op1(var2, var1)
op2_out = op2(op1_out, var2)
op3_out = op3(op2_out, var2, var2)
fg = FunctionGraph([var1, var2], [op1_out, op3_out], clone=False)
# Remove an output
fg.remove_output(1)
fg.check_integrity()
assert fg.outputs == [op1_out]
assert op3_out not in fg.clients
assert not any(op3_out.owner in clients
for clients in sum(fg.clients.values(), []))
# Remove an input
fg.remove_input(0)
fg.check_integrity()
assert var1 not in fg.variables
assert fg.inputs == [var2]
assert fg.outputs == []
assert not any(op1_out.owner in clients
for clients in sum(fg.clients.values(), []))
def test_remove_input(self):
var0 = MyVariable("var0")
var1 = MyVariable("var1")
var2 = MyVariable("var2")
var3 = MyVariable("var3")
var4 = MyVariable("var4")
op1_out = op1(var1, var0)
out0 = op2(op1_out, var2)
out1 = op1(var3, var4)
out1.name = "out1"
out2 = op1(out1, var0)
out2.name = "out2"
out3 = out1
fg = FunctionGraph(
[var0, var1, var2, var3, var4],
[out0, out1, out2, out3],
clone=False,
)
fg.remove_input(4)
fg.check_integrity()
assert fg.inputs == [var0, var1, var2, var3]
assert fg.outputs == [out0]
def test_remove_duplicates(self):
var1 = MyVariable("var1")
var2 = MyVariable("var2")
op1_out = op1(var2, var1)
op2_out = op2(op1_out, var2)
op3_out = op3(op2_out, var2, var2)
fg = FunctionGraph([var1, var1, var2], [op1_out, op3_out, op3_out],
clone=False)
fg.remove_output(2)
fg.check_integrity()
assert fg.outputs == [op1_out, op3_out]
fg.remove_input(0)
fg.check_integrity()
assert var1 not in fg.variables
assert fg.inputs == [var1, var2]
assert fg.outputs == []
