wehub-resource-sync
235 lines
8.7 KiB
Python
235 lines
8.7 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""The visitor and mutator infra for nn.Module."""
|
|
|
|
from typing import Any
|
|
|
|
from . import core as nn
|
|
|
|
|
|
class Mutator:
|
|
"""The mutator for nn.Module transform. Users can override the `visit_*` methods
|
|
to apply transform in different structures, or even override the `visit` method
|
|
to change the logic of traversal."""
|
|
|
|
def visit_module(self, name: str, node: nn.Module) -> Any:
|
|
"""The base visiting method for mutation of nn.Module nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.Module
|
|
The current node of nn.Module to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_effect(self, name: str, node: nn.Parameter) -> Any:
|
|
"""The base visiting method for mutation of nn.Parameter nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.Parameter
|
|
The current node of nn.Parameter to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_param(self, name: str, node: nn.Effect) -> Any:
|
|
"""The base visiting method for mutation of nn.Effect nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.Effect
|
|
The current node of nn.Effect to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_moduledict(self, name: str, node: nn.ModuleDict) -> Any:
|
|
"""The base visiting method for mutation of nn.ModuleDict nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.ModuleDict
|
|
The current node of nn.ModuleDict to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_modulelist(self, name: str, node: nn.ModuleList) -> Any:
|
|
"""The base visiting method for mutation of nn.ModuleList nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.ModuleList
|
|
The current node of nn.ModuleList to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_parameterdict(self, name: str, node: nn.ParameterDict) -> Any:
|
|
"""The base visiting method for mutation of nn.ParameterDict nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.ParameterDict
|
|
The current node of nn.ParameterDict to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit_parameterlist(self, name: str, node: nn.ParameterList) -> Any:
|
|
"""The base visiting method for mutation of nn.ParameterList nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : nn.ParameterList
|
|
The current node of nn.ParameterList to mutate.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
return self.visit(name, node)
|
|
|
|
def visit(self, name: str, node: Any) -> Any:
|
|
"""The base dispatching method for visiting of all nodes.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The name of the current node in parent's attribute.
|
|
|
|
node : Any
|
|
The current node to visit.
|
|
|
|
Returns
|
|
------
|
|
ret_node: Any
|
|
The new node to replace current node.
|
|
"""
|
|
|
|
def _get_child_name(parent: str, child: str) -> str:
|
|
"""Get the name of the child node/key given the parent's name."""
|
|
if parent == "":
|
|
# in the top level of the module
|
|
return child
|
|
else:
|
|
return f"{parent}.{child}"
|
|
|
|
if isinstance(node, nn.ParameterList):
|
|
for i in range(len(node)):
|
|
node[i] = self.visit_param(_get_child_name(name, str(i)), node[i])
|
|
elif isinstance(node, nn.ParameterDict):
|
|
for k, v in node.items():
|
|
node[k] = self.visit_param(_get_child_name(name, k), v)
|
|
elif isinstance(node, nn.ModuleList):
|
|
for i in range(len(node)):
|
|
if isinstance(node[i], nn.ParameterDict):
|
|
node[i] = self.visit_parameterdict(_get_child_name(name, str(i)), node[i])
|
|
elif isinstance(node[i], nn.ParameterList):
|
|
node[i] = self.visit_parameterlist(_get_child_name(name, str(i)), node[i])
|
|
elif isinstance(node[i], nn.ModuleDict):
|
|
node[i] = self.visit_moduledict(f"{name}.{i}", node[i])
|
|
elif isinstance(node[i], nn.ModuleList):
|
|
node[i] = self.visit_modulelist(f"{name}.{i}", node[i])
|
|
elif isinstance(node[i], nn.Module):
|
|
node[i] = self.visit_module(f"{name}.{i}", node[i])
|
|
elif isinstance(node[i], nn.Effect):
|
|
node[i] = self.visit_effect(f"{name}.{i}", node[i])
|
|
elif isinstance(node[i], nn.Parameter):
|
|
node[i] = self.visit_param(f"{name}.{i}", node[i])
|
|
elif isinstance(node, nn.ModuleDict):
|
|
for k, v in node.items():
|
|
if isinstance(v, nn.ParameterDict):
|
|
node[k] = self.visit_parameterdict(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.ParameterList):
|
|
node[k] = self.visit_parameterlist(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.ModuleDict):
|
|
node[k] = self.visit_moduledict(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.ModuleList):
|
|
node[k] = self.visit_modulelist(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.Module):
|
|
node[k] = self.visit_module(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.Effect):
|
|
node[k] = self.visit_effect(_get_child_name(name, k), v)
|
|
elif isinstance(v, nn.Parameter):
|
|
node[k] = self.visit_param(_get_child_name(name, k), v)
|
|
else:
|
|
for key, value in node.__dict__.items():
|
|
if isinstance(value, nn.ParameterDict):
|
|
setattr(node, key, self.visit_parameterdict(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.ParameterList):
|
|
setattr(node, key, self.visit_parameterlist(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.ModuleDict):
|
|
setattr(node, key, self.visit_moduledict(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.ModuleList):
|
|
setattr(node, key, self.visit_modulelist(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.Module):
|
|
setattr(node, key, self.visit_module(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.Effect):
|
|
setattr(node, key, self.visit_effect(_get_child_name(name, key), value))
|
|
elif isinstance(value, nn.Parameter):
|
|
setattr(node, key, self.visit_param(_get_child_name(name, key), value))
|
|
return node
|