!32291 [master]Bugfix for graph with tuple inputs

Merge pull request !32291 from caifubi/master-graph-tuple-input
This commit is contained in:
i-robot 2022-03-31 08:17:08 +00:00 committed by Gitee
commit aa6981ffa6
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
2 changed files with 34 additions and 0 deletions

View File

@ -460,6 +460,7 @@ GraphId GraphCompiler::CompileGraph(const FuncGraphPtr &func_graph, const Device
// The graph common optimization.
opt::BackendCommonOptimization(root_graph);
root_graph->SetInputNodes();
auto graph_id = CompileGraphImpl(root_graph, device_context);

View File

@ -383,3 +383,36 @@ def test_pynative_ms_function_with_dynamic_shape():
x = Tensor([[1, 1, 2], [3, 3, 5]], ms.int32)
output = test(x)
assert (output[0].asnumpy() == np.array([1, 2, 3, 5])).all()
@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.platform_x86_gpu_training
@pytest.mark.env_onecard
def test_pynative_ms_function_with_tuple_inputs():
"""
Feature: PyNative ms_function.
Description: PyNative ms_function with tuple inputs.
Expectation: The calculation result is correct.
"""
class Net(nn.Cell):
def __init__(self):
super(Net, self).__init__()
self.enable_tuple_broaden = True
@ms_function
def construct(self, grads):
new_grads = []
for grad in grads:
new_grads.append(grad + 1)
return new_grads
x = Tensor(np.ones([2, 2]), dtype=ms.int32)
y = Tensor(np.ones([2, 2]), dtype=ms.int32)
net = Net()
out = net((x, y))
assert (out[0].asnumpy() == np.ones([2, 2]) + 1).all()