mindspore.nn.Vjp

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

Computes the dot product between a vector v and the Jacobian of the given fn at the point given by the inputs.

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. Must be a tuple or a list.

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

Outputs:

A tuple with 2 Tensors or Tuple of Tensors:

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

  • vjp (Tensors or Tuple of Tensors) - The result of the dot product.

Supported Platforms:

Ascend GPU CPU

Examples

>>> from mindspore.nn import Vjp
>>> 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 = Vjp(Net())(x, y, v)
>>> print(output[0])
[[ 2. 10.]
 [30. 68.]]
>>> print(output[1][0])
[[ 3. 12.]
 [27. 48.]]
>>> print(output[1][1])
[[1. 1.]
 [1. 1.]]