mindspore.jvp

mindspore.jvp(fn, inputs, v, has_aux=False)

计算给定网络的雅可比向量积(Jacobian-vector product, JVP)。JVP的计算过程可参考 前向模式自动微分。

参数：

• fn (Union[Function, Cell]) - 待求导的函数或网络。以Tensor为入参，返回Tensor或Tensor数组。

• inputs (Union[Tensor, tuple[Tensor], list[Tensor]]) - 输入网络 fn 的入参。

• v (Union[Tensor, tuple[Tensor], list[Tensor]]) - 与雅可比矩阵相乘的向量，shape和type与网络的输入一致。

• has_aux (bool) - 若 has_aux 为 True ，只有 fn 的第一个输出参与 fn 的求导，其他输出将直接返回。此时， fn 的输出数量必须超过一个。默认值： False 。

返回：

• net_output (Union[Tensor, tuple[Tensor]]) - 输入网络的正向计算结果。

• jvp (Union[Tensor, tuple[Tensor]]) - 雅可比向量积的结果。

• aux_value (Union[Tensor, tuple[Tensor]], 可选) - 若 has_aux 为 True ，才返回 aux_value 。aux_value 是 fn(inputs) 的第一个除外的其他输出，且不参与 fn 的求导。

异常：

• TypeError - inputs 或 v 类型不符合要求。

支持平台：

Ascend GPU CPU

样例：

>>> import numpy as np
>>> from mindspore import jvp
>>> from mindspore import Tensor
>>> import mindspore.nn as nn
>>> class Net(nn.Cell):
...     def construct(self, x, y):
...         return x**3 + y
>>> x = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
>>> y = Tensor(np.array([[1, 2], [3, 4]]).astype(np.float32))
>>> v = Tensor(np.array([[1, 1], [1, 1]]).astype(np.float32))
>>> output = jvp(Net(), (x, y), (v, v))
>>> print(output[0])
[[ 2. 10.]
 [30. 68.]]
>>> print(output[1])
[[ 4. 13.]
 [28. 49.]]
>>>
>>> def fn(x, y):
...     return x ** 3 + y, y
>>> output, jvp_out, aux = jvp(fn, (x, y), (v, v), has_aux=True)
>>> print(output)
[[ 2. 10.]
 [30. 68.]]
>>> print(jvp_out)
[[ 4. 13.]
 [28. 49.]]
>>> print(aux)
[[ 1. 2.]
 [3. 4.]]