def _dfs(
    tensor: Tensor, op_type: str, visited: OrderedDict[Tensor, Set[Operator]]
) -> Set[Operator]:
    """Dfs pass to traverse the graph and collects descendant ops with type == op_type
    for every tensor, which is saved in `visited`.
    """
    if tensor in visited:
        return visited[tensor]

    # establish topological order in visited
    visited[tensor] = set()

    descendants = set()
    for op in tensor.dst_ops():
        outputs = op._attrs["outputs"]
        for output in outputs:
            descendants.update(_dfs(output, op_type, visited))

    # visited[tensor] should only contain descendants, not self
    visited[tensor].update(descendants)

    src_ops = list(tensor.src_ops())
    assert (
        len(src_ops) <= 1
    ), f"A tensor can't have more than 1 src_op in this stage, len(src_ops): {len(src_ops)}"

    src_op = src_ops[0] if len(src_ops) == 1 else None
    if src_op and src_op._attrs["op"] == op_type:
        descendants.add(src_op)
    return descendants
Dfs pass to traverse the graph and collects descendant ops with type == op_type
for every tensor, which is saved in `visited`.
python/aitemplate/compiler/transform/fuse_group_ops.py