def test_simple_transpose_squash():
    """
    Test eliminate consecutive transpose can be canceled
    """

    @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))])
    def prog(x0):
        x1 = mb.transpose(x=x0, perm=[1, 0])
        x2 = mb.transpose(x=x1, perm=[1, 0])
        x3 = mb.log(x=x2)
        x4 = mb.transpose(x=x3, perm=[1, 0])
        x5 = mb.transpose(x=x4, perm=[1, 0])
        x6 = mb.transpose(x=x5, perm=[1, 0])
        x7 = mb.transpose(x=x6, perm=[1, 0])
        return x7

    print("before:\n{}".format(prog))
    block = prog.functions["main"]
    assert len(block.find_ops(op_type="transpose")) == 6

    def can_squash(trans1, trans2):
        return (
            len(trans1.outputs) == 1
            and len(trans2.outputs) == 1
            and all(trans1.perm.val == trans2.perm.val)
        )

    # Find all candidate pairs transposes
    # we ignore all const (transpose_perm_%x), and add pairs of transpose op as
    # candidate. This won't generalize to more complicated program with other
    # shape invariant ops in between.
    candidates = []
    non_const_ops = [op for op in block.operations if op.op_type != "const"]
    for i in range(len(non_const_ops) - 1):
        op = non_const_ops[i]
        if len(candidates) and op == candidates[-1][1]:
            # op is already a squash candidate
            continue
        next_op = non_const_ops[i + 1]
        if (
            op.op_type == "transpose"
            and next_op.op_type == "transpose"
            and can_squash(op, next_op)
        ):
            candidates.append((op, next_op))

    # Remove each candidate pairs
    for (trans1, trans2) in candidates:
        before = trans1.inputs["x"]
        after = trans2.outputs[0]
        block.replace_uses_of_var_after_op(
            anchor_op=trans2, old_var=after, new_var=before
        )
        block.remove_ops([trans1, trans2])

    print("after:\n{}".format(prog))
    assert len(block.find_ops(op_type="transpose")) == 0

    Test eliminate consecutive transpose can be canceled
    
coremltools/converters/mil/mil/tests/test_block.py