def test_remove_op():
    """
    Test remove all ops and return empty program
    """

    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
    def prog(x0):
        x1 = mb.log(x=x0)
        return x1

    print("before:\n{}".format(prog))
    block = prog.functions["main"]
    assert len(block.operations) == 2
    x0 = block.inputs["x0"]
    ops = block.find_ops(op_type="log")
    block.set_outputs([x0])
    block.remove_ops(ops)
    print("after:\n{}".format(prog))
    assert len(block.operations) == 1
    assert len(block.inputs) == 1
    assert len(block.outputs) == 1
    assert block.inputs["x0"] == block.outputs[0]

    Test remove all ops and return empty program
    
coremltools/converters/mil/mil/tests/test_block.py