• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2021 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 jvp in graph mode"""
16import numpy as np
17import pytest
18import mindspore.nn as nn
19import mindspore.context as context
20from mindspore import Tensor
21from mindspore.nn.grad import Vjp
22
23
24class SingleInputNet(nn.Cell):
25    def construct(self, x):
26        return x**3
27
28
29class MultipleInputsOutputNet(nn.Cell):
30    def construct(self, x, y):
31        return 2*x, y**3
32
33
34@pytest.mark.level1
35@pytest.mark.platform_x86_cpu
36@pytest.mark.env_onecard
37@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
38def test_vjp_single_input_graph(mode):
39    """
40    Features: Class Vjp.
41    Description: Test whenther Vjp can calculate backward-mode diff correctly.
42    Expectation: No exception.
43    """
44    context.set_context(mode=mode)
45    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
46    v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
47    net = SingleInputNet()
48    expect_primal = Tensor(np.array([[1, 8], [27, 64]]).astype(np.float32))
49    expect_grad = Tensor(np.array([[3, 12], [27, 48]]).astype(np.float32))
50    primal, grad = Vjp(net)(x, v)
51    assert np.allclose(primal.asnumpy(), expect_primal.asnumpy())
52    assert np.allclose(grad.asnumpy(), expect_grad.asnumpy())
53
54
55
56@pytest.mark.level1
57@pytest.mark.platform_x86_cpu
58@pytest.mark.env_onecard
59@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
60def test_vjp_multiple_inputs_default_v_graph(mode):
61    """
62    Features: Class Vjp.
63    Description: Test whenther Vjp can calculate backward-mode diff correctly.
64    Expectation: No exception.
65    """
66    context.set_context(mode=mode)
67    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
68    y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
69    v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
70    net = MultipleInputsOutputNet()
71    expect_primal_0 = Tensor(np.array([[2, 4], [6, 8]]).astype(np.float32))
72    expect_primal_1 = Tensor(np.array([[1, 8], [27, 64]]).astype(np.float32))
73    expect_grad_0 = Tensor(np.array([[2, 2], [2, 2]]).astype(np.float32))
74    expect_grad_1 = Tensor(np.array([[3, 12], [27, 48]]).astype(np.float32))
75    primal, grad = Vjp(net)(x, y, (v, v))
76    assert isinstance(primal, tuple)
77    assert len(primal) == 2
78    assert np.allclose(primal[0].asnumpy(), expect_primal_0.asnumpy())
79    assert np.allclose(primal[1].asnumpy(), expect_primal_1.asnumpy())
80    assert isinstance(grad, tuple)
81    assert len(grad) == 2
82    assert np.allclose(grad[0].asnumpy(), expect_grad_0.asnumpy())
83    assert np.allclose(grad[1].asnumpy(), expect_grad_1.asnumpy())
84