def _get_stack(self, tfssa, root="main"):
        """
        Overwrite TFConverter._get_stack() as TF2 generates different sub-graphs.
        """

        # We're trying to get a order of how to loop through the graphs.
        # This is NOT necessarily a DAG.
        dep = {x: [] for x in tfssa.functions}
        for fname in tfssa.functions:
            for node in tfssa.functions[fname].graph.values():
                func_x, func_y = None, None

                if node.op in {"StatelessIf", "If"}:
                    func_x = node.attr.get("then_branch")
                    func_y = node.attr.get("else_branch")
                elif node.op in {"StatelessWhile", "While"}:
                    func_x = node.attr.get("body")
                    func_y = node.attr.get("cond")

                if func_x and fname not in dep[func_x]:
                    dep[func_x].append(fname)
                if func_y and fname not in dep[func_y]:
                    dep[func_y].append(fname)

        assert len(dep[root]) == 0
        graph_stack = simple_topsort(dep)

        return graph_stack

        Overwrite TFConverter._get_stack() as TF2 generates different sub-graphs.
        
coremltools/converters/mil/frontend/tensorflow2/converter.py