chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,234 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user