• 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 nn ops """
16import numpy as np
17import mindspore.nn as nn
18import mindspore.common.dtype as mstype
19
20from mindspore import Tensor
21from mindspore.ops import operations as P
22from mindspore import context
23
24context.set_context(mode=context.GRAPH_MODE)
25
26
27def test_cast_op_attr():
28    class CastNet(nn.Cell):
29        def __init__(self):
30            super(CastNet, self).__init__()
31            self.cast = P.Cast()
32
33        def construct(self, x, t):
34            return self.cast(x, t)
35
36    class CastTypeTest(nn.Cell):
37        def __init__(self, net):
38            super(CastTypeTest, self).__init__()
39            self.net = net
40            self.cast = P.Cast()
41
42        def construct(self, x, y, z):
43            cast_op = self.cast
44            t1 = cast_op(x, mstype.float32)
45            t2 = cast_op(y, mstype.int32)
46            cast_net = self.net
47            t3 = cast_net(x, mstype.float16)
48            t4 = cast_net(y, mstype.int32)
49            t5 = cast_net(z, mstype.float16)
50            return (t1, t2, t3, t4, t5)
51
52    net = CastTypeTest(CastNet())
53    t1 = Tensor(np.ones([1, 16, 1, 1918]).astype(np.int32))
54    t2 = Tensor(np.ones([1, 16, 1, 3840]).astype(np.float32))
55    t3 = Tensor(np.ones([1, 16, 1, 1918]).astype(np.int32))
56    out = net(t1, t2, t3)
57    assert out[0].asnumpy().dtype == np.float32
58    assert out[1].asnumpy().dtype == np.int32
59    assert out[2].asnumpy().dtype == np.float16
60    assert out[3].asnumpy().dtype == np.int32
61    assert out[4].asnumpy().dtype == np.float16
62