mindspore.nn.Jvp

class mindspore.nn.Jvp(fn)[source]

Compute the jacobian-vector-product of the given fn. Jvp is equivalent to forward mode autodiff.

Parameters

fn (Cell) – The fn that takes Tensor inputs and returns a tuple of Tensors or a Tensor.

Inputs:

  • inputs (Tensors) - The inputs to fn.

  • v (Tensors or Tuple of Tensors) - The vector for which the Jacobian vector product is computed. Must have the same size as the input of fn.

Outputs:

A tuple with 2 Tensors or Tuple of Tensors:

  • net_output (Tensors or Tuple of Tensors) - The output of fn(inputs).

  • jvp (Tensors or Tuple of Tensors) - The result of the jacobian vector product.

Supported Platforms:

Ascend GPU CPU

Examples

>>> from mindspore.nn import Jvp
>>> 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.]]