def test_substitute_nested_op():
    """"
    Replace an conditional op with nested block
    """

    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4)), mb.TensorSpec(shape=(2, 4))])
    def prog(x0, y0):
        pred = mb.less(x=x0, y=y0)
        z = mb.cond(
            pred=pred, _true_fn=lambda: mb.abs(x=x0), _false_fn=lambda: mb.abs(x=y0)
        )
        z1 = mb.log(x=z)
        return z1

    print("before:\n{}".format(prog))
    block = prog.functions["main"]
    assert len(block.find_ops(op_type="less")) == 1
    assert len(block.find_ops(op_type="abs")) == 2
    assert len(block.find_ops(op_type="cond")) == 1
    assert len(block.find_ops(op_type="log")) == 1

    cond = block.find_ops(op_type="cond")[0]
    x0 = block.inputs["x0"]
    z = cond.outputs[0]
    block.replace_uses_of_var_after_op(anchor_op=None, old_var=z, new_var=x0)

    # removing cond will also remove the abs ops within its block
    block.remove_ops([cond])

    print("after:\n{}".format(prog))
    assert len(block.find_ops(op_type="less")) == 1
    assert len(block.find_ops(op_type="log")) == 1
    assert len(block.find_ops(op_type="cond")) == 0
    assert len(block.find_ops(op_type="abs")) == 0
"
    Replace an conditional op with nested block
    
coremltools/converters/mil/mil/tests/test_block.py