def _remove_no_op_expands(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    Remove no-op expands from the graph. A no-op expand is one
    that doesn't expand any singleton dimensions to values greater
    than one.

    x = Tensor([1, 2, 3])
    y1 = ops.expand()(x, [-1, -1, -1])  # no-op
    y2 = ops.expand()(x, [1, 2, -1])  # no-op
    """
    ops = graph_utils.get_sorted_ops(sorted_graph)
    for op in ops:
        if op._attrs["op"] != "expand":
            continue

        outputs = op._attrs["outputs"]
        assert len(outputs) == 1, "expand must only have 1 output"
        expand_output = outputs[0]

        if expand_output._attrs["is_output"]:
            continue

        inputs = op._attrs["inputs"]
        assert len(inputs) >= 1, "expand must have at least 1 input"
        expand_input = inputs[0]

        assert len(op._attrs["dim_types"]) == len(
            expand_output._attrs["shape"]
        ), "expand must have dim_type for every output dimension"

        # If we just keep every dimension as-is, it is a no-op
        if any(dt != ExpandDimensionType.KEEP_DIM for dt in op._attrs["dim_types"]):
            continue

        # This expand is a no-op, so we know that these shapes should
        # be the same. However, the shape inference system may not be aware
        # of that due to different IntVar names.
        expand_input._attrs["shape"] = expand_output._attrs["shape"]
        for dst in list(expand_output.dst_ops()):
            transform_utils.replace_tensor_for_op(dst, expand_output, expand_input)

        transform_utils.remove_tensor_from_sorted_graph(expand_output)

    return transform_utils.sanitize_sorted_graph(sorted_graph)
Remove no-op expands from the graph. A no-op expand is one
that doesn't expand any singleton dimensions to values greater
than one.

x = Tensor([1, 2, 3])
y1 = ops.expand()(x, [-1, -1, -1])  # no-op
y2 = ops.expand()(x, [1, 2, -1])  # no-op
python/aitemplate/compiler/transform/remove_no_ops.py