def rewrite(self, debug: bool = False, custom_registry: _BaseRegistry = None):
        """
        Method which goes through each node in the graph and then each transform
        in the registry. If the transform pattern matches against the node, will
        compute the replacement module and replace the node with this module in
        the graph. This method then goes through the modified graph to delete
        unused operators and dead code. Finally, it recompiles and sanity-checks the
        graph before preparing the rewritten model.
        There is a debug flag which enables the printing of detailed information about
        the rewritten model and modified graph.
        """

        modules = dict(self.model.named_modules())
        registry = (
            _TransformRegistry.get_registry_values()
            if custom_registry is None
            else custom_registry
        )

        # fx represents its graph as an ordered list of
        # nodes so we can iterate through them.
        for node in self.graph_module.graph.nodes:
            for transform in registry:
                if transform.match_pattern(node, modules, self.env):
                    repl = transform.get_replacement(
                        self.model, node, self.shapes[node.name], modules, self.env
                    )
                    if transform.get_replacement_type() is _ReplacementType.kMODULE:
                        self._replace_node_module(node, repl)
                    elif transform.get_replacement_type() is _ReplacementType.kFUNCTION:
                        self._replace_node_function(node, repl[0], repl[1], repl[2])
                    else:
                        raise Exception("Replacement type is not supported")
                    # Regenerate modules
                    modules = dict(self.model.named_modules())

        self.graph_module.delete_all_unused_submodules()
        self.graph_module.graph.eliminate_dead_code()
        self.graph_module.recompile()
        # Does some checks to make sure the graph is well-formed.
        self.graph_module.graph.lint()
        self.model = _torch.fx.GraphModule(self.model, self.graph_module.graph)

        if debug:
            print("=======================================================")
            print("Printing human-readable graph module")
            self.graph_module.print_readable()
            print("=======================================================")
            print("Printing graph from graph module")
            print(self.graph_module.graph)
            print("=======================================================")
            print("Printing tabular graph")
            self.graph_module.graph.print_tabular()
            print("=======================================================")
            print("Printing named modules")
            for name, _ in self.model.named_modules():
                print("Found module with name", name)
            print("=======================================================")

        Method which goes through each node in the graph and then each transform
        in the registry. If the transform pattern matches against the node, will
        compute the replacement module and replace the node with this module in
        the graph. This method then goes through the modified graph to delete
        unused operators and dead code. Finally, it recompiles and sanity-checks the
        graph before preparing the rewritten model.
        There is a debug flag which enables the printing of detailed information about
        the rewritten model and modified graph.
        
coremltools/optimize/torch/_utils/graph_utils.py