def _partition_subgraphs(ops: Set[Operator]) -> Dict[str, Set[Operator]]:
    """
    Given ops of candidate graph of fused_elementwise op graph and partition
    into subgraph based on output shape, returns dict of
    {output shape: ops to form subgraph based on the shape}
    """
    # Partition graph of elementwise into subgraph based on output shape.
    output_op_map = collections.defaultdict(set)
    for op in ops:
        shapes = []
        # Find output nodes
        for output_tensor in op._attrs["outputs"]:
            if (
                output_tensor._attrs["is_output"]
                or len(output_tensor._attrs["dst_ops"] - ops) > 0
            ):
                shapes.append("_".join(map(str, output_tensor._attrs["shape"])))
        # Find anscestor of output node.
        # Outputs with the same shape should form the same graph
        if shapes:
            key = "|".join(shapes)
            op_set = output_op_map[key]
            for anc_op in ops:
                if transform_utils.is_ancestor(anc_op, op):
                    op_set.add(anc_op)
            op_set.add(op)
    return output_op_map

    Given ops of candidate graph of fused_elementwise op graph and partition
    into subgraph based on output shape, returns dict of
    {output shape: ops to form subgraph based on the shape}
    
python/aitemplate/compiler/transform/fuse_ops.py