forked from huawei/mindspore2022
128 lines
5.5 KiB
Python
128 lines
5.5 KiB
Python
# Copyright 2022 Huawei Technologies Co., Ltd
|
|
#
|
|
# Licensed 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.
|
|
# ============================================================================
|
|
import ast
|
|
|
|
from mindspore.nn import Cell
|
|
from mindspore.rewrite import ScopedValue
|
|
from mindspore.rewrite.node import Node
|
|
|
|
|
|
class FakeCell(Cell):
|
|
def construct(self, input1, input2, cool_boy=None):
|
|
return input1 + input2 + cool_boy
|
|
|
|
|
|
class FakeCell2(Cell):
|
|
def construct(self, a, b, d, e, *args, f=6, **kwargs):
|
|
return a + b + d + e + sum(args) + f + sum(kwargs.values())
|
|
|
|
|
|
class FakeCell3(Cell):
|
|
def construct(self, a, b, *args, f=6, h=7, **kwargs):
|
|
return a + b + f + h + sum(args) + sum(kwargs.values())
|
|
|
|
|
|
def test_create_by_cell():
|
|
"""
|
|
Feature: Python api create_call_buildin_op of Node of Rewrite.
|
|
Description: Call create_call_buildin_op to create a CallCell node.
|
|
Expectation: Success.
|
|
"""
|
|
node = Node.create_call_buildin_op(FakeCell(), None, ['x'], 'new_conv',
|
|
[ScopedValue.create_naming_value('x'), ScopedValue.create_variable_value(1)],
|
|
{"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv')
|
|
assert node._args_num == 2
|
|
assert node._kwargs_num == 1
|
|
assert node._normalized_args_keys == ["input1", "input2", "cool_boy"]
|
|
|
|
assert node._normalized_args == {
|
|
"input1": ScopedValue.create_naming_value('x'),
|
|
"input2": ScopedValue.create_variable_value(1),
|
|
"cool_boy": ScopedValue.create_naming_value('Naroto')
|
|
}
|
|
|
|
ast_node: ast.Assign = node.get_ast()
|
|
assign_value: ast.Call = ast_node.value
|
|
args_ast = assign_value.args
|
|
keywords_ast = assign_value.keywords
|
|
assert len(args_ast) == 2
|
|
assert len(keywords_ast) == 1
|
|
assert keywords_ast[0].arg == "cool_boy"
|
|
assert isinstance(args_ast[0], ast.Name)
|
|
assert args_ast[0].id == "x"
|
|
assert isinstance(args_ast[1], ast.Constant)
|
|
assert args_ast[1].value == 1
|
|
keyword_value_3 = keywords_ast[0].value
|
|
assert isinstance(keyword_value_3, ast.Name)
|
|
assert keyword_value_3.id == "Naroto"
|
|
|
|
node.set_arg(ScopedValue.create_variable_value(2), 1)
|
|
assert isinstance(node.get_normalized_args().get("input2"), ScopedValue)
|
|
assert node.get_normalized_args().get("input2").value == 2
|
|
ast_node: ast.Assign = node.get_ast()
|
|
assign_value: ast.Call = ast_node.value
|
|
args_ast = assign_value.args
|
|
assert args_ast[1].value == 2
|
|
|
|
args = node.get_args()
|
|
assert args == [ScopedValue.create_naming_value('x'), ScopedValue.create_variable_value(2)]
|
|
kwargs = node.get_kwargs()
|
|
assert kwargs == {"cool_boy": ScopedValue.create_naming_value('Naroto')}
|
|
|
|
|
|
def test_create_by_cell2():
|
|
"""
|
|
Feature: Python api create_call_buildin_op of Node of Rewrite.
|
|
Description: Call create_call_buildin_op to create a CallCell node.
|
|
Expectation: Success.
|
|
"""
|
|
node = Node.create_call_buildin_op(FakeCell2(), None, ['x'], 'new_conv',
|
|
[ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"),
|
|
ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"),
|
|
ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")],
|
|
{"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv')
|
|
assert node.get_normalized_args() == {
|
|
"a": ScopedValue.create_naming_value('x'),
|
|
"b": ScopedValue.create_naming_value('x'),
|
|
"d": ScopedValue.create_naming_value('x'),
|
|
"e": ScopedValue.create_naming_value('x'),
|
|
"args_4": ScopedValue.create_naming_value('x'),
|
|
"args_5": ScopedValue.create_naming_value('x'),
|
|
"cool_boy": ScopedValue.create_naming_value('Naroto'),
|
|
}
|
|
|
|
|
|
def test_create_by_cell3():
|
|
"""
|
|
Feature: Python api create_call_buildin_op of Node of Rewrite.
|
|
Description: Call create_call_buildin_op to create a CallCell node.
|
|
Expectation: Success.
|
|
"""
|
|
node = Node.create_call_buildin_op(FakeCell3(), None, ['x'], 'new_conv',
|
|
[ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x"),
|
|
ScopedValue.create_naming_value('x'), ScopedValue.create_naming_value("x")],
|
|
{"h": ScopedValue.create_naming_value(1),
|
|
"f": ScopedValue.create_naming_value(2),
|
|
"cool_boy": ScopedValue.create_naming_value('Naroto')}, 'new_conv')
|
|
assert node.get_normalized_args() == {
|
|
"a": ScopedValue.create_naming_value('x'),
|
|
"b": ScopedValue.create_naming_value('x'),
|
|
"args_2": ScopedValue.create_naming_value('x'),
|
|
"args_3": ScopedValue.create_naming_value('x'),
|
|
"f": ScopedValue.create_naming_value(2),
|
|
"h": ScopedValue.create_naming_value(1),
|
|
"cool_boy": ScopedValue.create_naming_value('Naroto'),
|
|
}
|