def _fuse_expand_elementwise(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    Eliminate expand ops that occur before elementwise when broadcasting
    in elementwise can handle the unexpanded input.

    Example:
    x = Tensor([1, 2, 3])
    y = Tensor([3, 2, 3])
    z = ops.elementwise(FuncEnum.ADD)(ops.expand()(x, [3, 2, 3]), y)

    The expand here is not required because elementwise broadcasting will just
    do the right thing.

    Note that this must occur before any pass that fuses elementwise into
    other ops.
    """

    def _is_compatible_with_broadcasting(
        expand_output_dim: IntVar, elementwise_input_dim: IntVar
    ) -> bool:
        return expand_output_dim == elementwise_input_dim or is_singleton_dimension(
            expand_output_dim
        )

    def _replace_jagged_int_var(shape: List[IntVar]):
        """
        If shape[0] is a JaggedIntVar, replace it with
        the corresponding maximum dense shape.
        """
        if shape and isinstance(shape[0], JaggedIntVar):
            return shape[0].get_max_dense_shape() + shape[1:]
        return shape

    for op in graph_utils.get_sorted_ops(sorted_graph):
        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

        expand_output_shape = _replace_jagged_int_var(expand_output._attrs["shape"])

        def _can_fuse_with(dst_op: Operator) -> bool:
            if dst_op._attrs["op"] != "elementwise":
                return False

            for elementwise_input in dst_op._attrs["inputs"]:
                if elementwise_input is expand_output:
                    continue

                elementwise_input_shape = _replace_jagged_int_var(
                    elementwise_input._attrs["shape"]
                )

                if not all(
                    _is_compatible_with_broadcasting(dim_a, dim_b)
                    for dim_a, dim_b in zip(
                        expand_output_shape,
                        elementwise_input_shape,
                    )
                ):
                    return False
            return True

        if not all(_can_fuse_with(dst) for dst in expand_output._attrs["dst_ops"]):
            continue

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

        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)
Eliminate expand ops that occur before elementwise when broadcasting
in elementwise can handle the unexpanded input.

Example:
x = Tensor([1, 2, 3])
y = Tensor([3, 2, 3])
z = ops.elementwise(FuncEnum.ADD)(ops.expand()(x, [3, 2, 3]), y)

The expand here is not required because elementwise broadcasting will just
do the right thing.

Note that this must occur before any pass that fuses elementwise into
other ops.
python/aitemplate/compiler/transform/remove_no_ops.py