def _fuse_layernorm_ops(
    op_group: List[Operator], sorted_graph: List[Tensor]
) -> List[Tensor]:
    """
    Replace a group of ops with a single group op
    """
    # Make the order deterministic
    # Sort by gamma name
    op_group.sort(key=lambda x: x._attrs["inputs"][1]._attrs["name"])

    # gather inputs
    num_inputs = len(op_group[0]._attrs["inputs"])
    group_inputs = [[] for _ in range(num_inputs)]
    normalized_shapes = []
    for op in op_group:
        normalized_shapes.append(op._attrs["normalized_shape"])
        for i, input in enumerate(op._attrs["inputs"]):
            group_inputs[i].append(input)

    # remove dst_ops from inputs
    for op in op_group:
        for input in op._attrs["inputs"]:
            transform_utils.remove_dst_op_from_tensor(input, op)

    # create group op
    op_type = op_group[0]._attrs["op"]
    group_op = (
        ops.group_layernorm
        if op_type == "layernorm"
        else ops.group_layernorm_sigmoid_mul
    )
    eps = op_group[0]._attrs["eps"]
    group_outputs = group_op()(
        group_inputs[0], group_inputs[1], group_inputs[2], normalized_shapes, eps
    )

    for i, op in enumerate(op_group):
        new_output = group_outputs[i]
        op_output = op._attrs["outputs"][0]
        transform_utils.replace_tensor(op_output, new_output)

    # sorted_graph is no longer sorted here
    sorted_graph.extend(group_outputs)
    return sorted_graph
Replace a group of ops with a single group op
python/aitemplate/compiler/transform/fuse_group_ops.py