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.]]