• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 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# ============================================================================
15""" test_hypermap_partial """
16import numpy as np
17
18import mindspore.common.dtype as mstype
19import mindspore.nn as nn
20from mindspore import Tensor, context
21from mindspore.common.api import ms_function
22from mindspore.ops import composite as C
23from mindspore.ops import functional as F
24from mindspore.ops import operations as P
25
26context.set_context(mode=context.GRAPH_MODE)
27
28
29def test_hypermap_specialize_param():
30    class Net(nn.Cell):
31        """ Net definition """
32
33        def __init__(self):
34            super(Net, self).__init__()
35            self.mul = P.Mul()
36
37        def construct(self, x, y):
38            ret = self.mul(x, y)
39            return ret
40
41    factor1 = Tensor(5, dtype=mstype.int32)
42    x = Tensor(np.ones([1]).astype(np.int32))
43    y = Tensor(np.ones([2]).astype(np.int32))
44    net = Net()
45    hypermap = C.HyperMap()
46
47    @ms_function
48    def hypermap_specialize_param():
49        ret1 = hypermap(F.partial(net, factor1), (x, y))
50        # List will be converted to Tuple in SimlifyDataStructurePass.
51        ret2 = hypermap(F.partial(net, factor1), [x, y])
52        return ret1, ret2
53
54    expected_ret = (Tensor(np.full(1, 5).astype(np.int32)), Tensor(np.full(2, 5).astype(np.int32)))
55    ret = hypermap_specialize_param()
56    assert ret[0][0].asnumpy() == expected_ret[0].asnumpy()
57    assert np.all(ret[0][1].asnumpy() == expected_ret[1].asnumpy())
58    assert ret[1][0].asnumpy() == list(expected_ret[0].asnumpy())
59    assert np.all(ret[1][1].asnumpy() == list(expected_ret[1].asnumpy()))
60