1# Copyright 2019 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15import numpy as np 16 17from mindspore import Tensor 18from mindspore.common.api import ms_function 19from mindspore.ops import operations as P 20 21 22def test_nest_range_transpose(): 23 batch_size = 2 24 num_layers = 5 25 batch_tuple = tuple(Tensor(np.array(np.ones((2, 3)) * 0.01)) for i in range(batch_size)) 26 layers_tuple = tuple(Tensor(np.array(np.ones((3, 4)) * 0.02)) for i in range(num_layers)) 27 transpose1 = P.Transpose() 28 29 @ms_function() 30 def invoke_range(): 31 out1 = () 32 for m in range(num_layers): 33 out1 += (transpose1(layers_tuple[m], (1, 0)),) 34 # Both for loop will the same range symbol as phi node, when range primitive is converted 35 # to DoSigature MetaFuncGraph, that MetaFuncGraph will take 2 and 5 as argument, so there is 36 # 2 entries in that MetaFuncGraphEvaluator, that will make Specialier try to use AnyValue to 37 # FindGeneralized for S-make_range MetaFuncGraph but it will fail as AnyValue is not constant. 38 for i in range(batch_size): 39 out1 += (transpose1(batch_tuple[i], (1, 0)),) 40 for j in range(num_layers): 41 out1 += (transpose1(layers_tuple[j], (1, 0)),) 42 return out1 43 44 print(invoke_range()) 45 46 47def test_nest_range_simple(): 48 batch_size = 2 49 num_layers = 5 50 batch_tuple = tuple(Tensor(np.array(np.ones((2, 3)) * 0.01)) for i in range(batch_size)) 51 layers_tuple = tuple(Tensor(np.array(np.ones((3, 4)) * 0.02)) for i in range(num_layers)) 52 53 @ms_function() 54 def invoke_range(): 55 out1 = () 56 for m in range(num_layers): 57 out1 += (layers_tuple[m],) 58 for i in range(batch_size): 59 out1 += (batch_tuple[i],) 60 for j in range(num_layers): 61 out1 += (layers_tuple[j],) 62 return out1 63 64 print(invoke_range()) 65