!26345 add examples of function grad, jvp and vjp

Merge pull request !26345 from chenzhuo/jvp

mindspore/ops/functional.py

@@ -195,6 +195,31 @@ def grad(fn, grad_position=0, sens_param=False):
 
     Returns:
         Function, returns the gradient function for the input function or cell.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> import mindspore.context as context
+        >>> from mindspore import Tensor
+        >>> from mindspore.ops.functional import grad
+        >>> context.set_context(mode=context.GRAPH_MODE)
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x*y*z
+        >>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+        >>> y = Tensor(np.array([[-2, 3], [-1, 2]]).astype(np.float32))
+        >>> z = Tensor(np.array([[0, 3], [5, -1]]).astype(np.float32))
+        >>> net = Net()
+        >>> output = grad(net, grad_position=(1, 2))(x, y, z)
+        >>> print(output)
+        (Tensor(shape=[2, 2], dtype=Float32, value=
+        [[ 0.00000000e+00,  6.00000000e+00],
+         [ 1.50000000e+01, -4.00000000e+00]]), Tensor(shape=[2, 2], dtype=Float32, value=
+        [[-2.00000000e+00,  6.00000000e+00],
+         [-3.00000000e+00,  8.00000000e+00]]))
     """
     grad_position = _convert_grad_position_type(grad_position)
     if sens_param:

tests/ut/python/nn/gradient/test_function_jvp_graph.py

@@ -15,6 +15,7 @@
 """test function jvp in graph mode"""
 
 import numpy as np
+import pytest
 import mindspore.nn as nn
 import mindspore.context as context
 from mindspore import Tensor
@@ -143,3 +144,16 @@ def test_jvp_multiple_inputs_multiple_outputs_custom_v_graph():
     v2 = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
     net = MultipleInputMultipleOutputNet()
     jvp(net, (x, y), (v1, v2))
+
+
+def test_jvp_wrong_input_type_graph():
+    """
+    Features: Function jvp
+    Description: Test jvp with wrong input type in graph mode.
+    Expectation: No exception.
+    """
+    x = 1
+    v = 1
+    net = SingleInputSingleOutputNet()
+    with pytest.raises(TypeError):
+        jvp(net, x, v)

tests/ut/python/nn/gradient/test_function_jvp_pynative.py

@@ -15,6 +15,7 @@
 """test function jvp in pynative mode """
 
 import numpy as np
+import pytest
 import mindspore.nn as nn
 import mindspore.context as context
 from mindspore import Tensor
@@ -142,3 +143,16 @@ def test_jvp_multiple_inputs_single_output_custom_v_pynative():
     v2 = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
     net = MultipleInputSingleOutputNet()
     jvp(net, (x, y), (v1, v2))
+
+
+def test_jvp_wrong_input_type_pynative():
+    """
+    Features: Function jvp
+    Description: Test jvp with wrong input type in pynative mode.
+    Expectation: No exception.
+    """
+    x = 1
+    v = 1
+    net = SingleInputSingleOutputNet()
+    with pytest.raises(TypeError):
+        jvp(net, x, v)

tests/ut/python/nn/gradient/test_function_vjp_graph.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """test vjp in graph mode"""
 import numpy as np
+import pytest
 import mindspore.nn as nn
 import mindspore.context as context
 from mindspore import Tensor
@@ -55,3 +56,16 @@ def test_vjp_multiple_inputs_default_v_graph():
     v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
     net = MultipleInputsOutputNet()
     vjp(net, (x, y), (v, v))
+
+
+def test_vjp_wrong_input_type_graph():
+    """
+    Features: Function vjp
+    Description: Test vjp with wrong input type in graph mode.
+    Expectation: No exception.
+    """
+    x = 1
+    v = 1
+    net = SingleInputNet()
+    with pytest.raises(TypeError):
+        vjp(net, x, v)

tests/ut/python/nn/gradient/test_function_vjp_pynative.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """test vjp in pynative mode"""
 import numpy as np
+import pytest
 import mindspore.nn as nn
 import mindspore.context as context
 from mindspore import Tensor
@@ -55,3 +56,16 @@ def test_vjp_multiple_inputs_default_v_pynative():
     v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
     net = MultipleInputsOutputNet()
     vjp(net, (x, y), (v, v))
+
+
+def test_vjp_wrong_input_type_pynative():
+    """
+    Features: Function vjp
+    Description: Test vjp with wrong input type in pynative mode.
+    Expectation: No exception.
+    """
+    x = 1
+    v = 1
+    net = SingleInputNet()
+    with pytest.raises(TypeError):
+        vjp(net, x, v)

tests/ut/python/nn/gradient/test_grad_graph.py

@@ -0,0 +1,73 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""test function grad in graph mode"""
+import numpy as np
+import mindspore.nn as nn
+import mindspore.context as context
+from mindspore import Tensor
+from mindspore.ops.functional import grad
+
+context.set_context(mode=context.GRAPH_MODE)
+
+
+class SingleInputSingleOutputNet(nn.Cell):
+    def construct(self, x):
+        return x**3
+
+
+class MultipleInputsMultipleOutputsNet(nn.Cell):
+    def construct(self, x, y, z):
+        return x**2 + y**2 + z**2, x*y*z
+
+
+def function(x, y, z):
+    return x**2 + y**2 + z**2, x*y*z
+
+
+def test_grad_single_input_single_output_cell_graph():
+    """
+    Features: Function grad.
+    Description: Test F.grad with single input and single output net in graph mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    net = SingleInputSingleOutputNet()
+    grad(net)(x)
+
+
+def test_grad_multiple_inputs_multiple_outputs_cell_graph():
+    """
+    Features: Function grad.
+    Description: Test F.grad with multiple inputs and multiple outputs net in graph mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    y = Tensor(np.array([[-2, 3], [-1, 2]]).astype(np.float32))
+    z = Tensor(np.array([[0, 3], [5, -1]]).astype(np.float32))
+    net = MultipleInputsMultipleOutputsNet()
+    grad(net, grad_position=(1, 2))(x, y, z)
+
+
+def test_grad_function_with_sens_graph():
+    """
+    Features: Function grad.
+    Description: Test F.grad with function setting sens_param in graph mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    y = Tensor(np.array([[-2, 3], [-1, 2]]).astype(np.float32))
+    z = Tensor(np.array([[0, 3], [5, -1]]).astype(np.float32))
+    v = Tensor(np.array([[-1, 3], [2, 1]]).astype(np.float32))
+    grad(function, grad_position=(1, 2), sens_param=True)(x, y, z, (v, v))

tests/ut/python/nn/gradient/test_grad_pynative.py

@@ -0,0 +1,73 @@
+# Copyright 2021 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""test function grad in pynative mode"""
+import numpy as np
+import mindspore.nn as nn
+import mindspore.context as context
+from mindspore import Tensor
+from mindspore.ops.functional import grad
+
+context.set_context(mode=context.PYNATIVE_MODE)
+
+
+class SingleInputSingleOutputNet(nn.Cell):
+    def construct(self, x):
+        return x**3
+
+
+class MultipleInputsMultipleOutputsNet(nn.Cell):
+    def construct(self, x, y, z):
+        return x**2 + y**2 + z**2, x*y*z
+
+
+def function(x, y, z):
+    return x**2 + y**2 + z**2, x*y*z
+
+
+def test_grad_single_input_single_output_cell_pynative():
+    """
+    Features: Function grad.
+    Description: Test F.grad with single input and single output net in pynative mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    net = SingleInputSingleOutputNet()
+    grad(net)(x)
+
+
+def test_grad_multiple_inputs_multiple_outputs_cell_pynative():
+    """
+    Features: Function grad.
+    Description: Test F.grad with multiple inputs and multiple outputs net in pynative mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    y = Tensor(np.array([[-2, 3], [-1, 2]]).astype(np.float32))
+    z = Tensor(np.array([[0, 3], [5, -1]]).astype(np.float32))
+    net = MultipleInputsMultipleOutputsNet()
+    grad(net, grad_position=(1, 2))(x, y, z)
+
+
+def test_grad_function_with_sens_pynative():
+    """
+    Features: Function grad.
+    Description: Test F.grad with function setting sens_param in pynative mode.
+    Expectation: No exception.
+    """
+    x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
+    y = Tensor(np.array([[-2, 3], [-1, 2]]).astype(np.float32))
+    z = Tensor(np.array([[0, 3], [5, -1]]).astype(np.float32))
+    v = Tensor(np.array([[-1, 3], [2, 1]]).astype(np.float32))
+    grad(function, grad_position=(1, 2), sens_param=True)(x, y, z, (v, v))