def _get_inputs_outputs(
    partitioned_ops: Set[Operator], all_ops: Set[Operator]
) -> List[List[Tensor]]:
    """
    Given ops of a partitioned subgraph based on output shape, and ops of full graph
    to form a complete graph with fused_elementwise op, returns all inputs/outputs of
    the ops and the external input/output of the subgraph, which will serve as input/output
    of fused_elementwise op.
    """
    external_inputs, external_outputs = [], []
    tmp_inputs, tmp_outputs = [], []

    for op in partitioned_ops:
        for input_tensor in op._attrs["inputs"]:
            tmp_inputs.append(input_tensor)
            src_ops = set(input_tensor._attrs["src_ops"])
            if (len(src_ops) == 0 or len(src_ops - all_ops) > 0) and (
                not input_tensor.is_a_const_num()
            ):
                external_inputs.append(input_tensor)
            assert op in input_tensor._attrs["dst_ops"]
        for output_tensor in op._attrs["outputs"]:
            tmp_outputs.append(output_tensor)
            dst_ops = set(output_tensor._attrs["dst_ops"])
            if output_tensor._attrs["is_output"] or len(dst_ops - all_ops) > 0:
                external_outputs.append(output_tensor)
            assert len(output_tensor._attrs["src_ops"]) == 1
            assert list(output_tensor._attrs["src_ops"])[0] == op

    # dict.fromkeys takes unique tensors and preserves the ordering.
    external_inputs = list(dict.fromkeys(external_inputs))
    external_outputs = list(dict.fromkeys(external_outputs))
    tmp_inputs = list(dict.fromkeys(tmp_inputs))
    tmp_outputs = list(dict.fromkeys(tmp_outputs))

    assert set(external_inputs) == set(tmp_inputs) - set(
        tmp_outputs
    ), "external_inputs: {} is not equal to tmp_inputs: {} - tmp_outputs: {}.".format(
        external_inputs, tmp_inputs, tmp_outputs
    )
    assert (
        len(set(tmp_outputs) - set(tmp_inputs) - set(external_outputs)) == 0
    ), "tmp_outputs: {} - tmp_inputs: {} - external_outputs: {} is not empty.".format(
        tmp_outputs, tmp_inputs, external_outputs
    )
    assert (
        len(set(external_outputs) - set(tmp_outputs)) == 0
    ), "external_outputs: {} - tmp_outputs: {} is not empty.".format(
        external_outputs, tmp_outputs
    )

    return [tmp_inputs, tmp_outputs, external_inputs, external_outputs]
Given ops of a partitioned subgraph based on output shape, and ops of full graph
to form a complete graph with fused_elementwise op, returns all inputs/outputs of
the ops and the external input/output of the subgraph, which will serve as input/output
of fused_elementwise op.
python/aitemplate/compiler/transform/fuse_ops.py