sciai.operators.JacobianWeights

class sciai.operators.JacobianWeights(model, out_shape, out_type=ms.float32)

Jacobian matrix with respect to weight(s). The last tensor in the input Tensor tuple is the weight Parameter, and the remainders are the inputs of network.

Parameters

• model (Cell) – Network for jacobian result with respect to weights.

• out_shape (tuple) – Output shape of the netword.

• out_type (type) – Mindspore data type. Default: ms.float32.

Inputs:

• x (tuple[Tensor]) - Tensors of the network input and the weight to find jacobian matrix.

Outputs:

Tensor, Jacobian matrix with respect to the given weights.

Examples

>>> import mindspore as ms
>>> from mindspore import nn, ops
>>> from sciai.operators import JacobianWeights
>>> class Net1In1OutTensor(nn.Cell):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.dense1 = nn.Dense(2, 1)
>>>     def construct(self, x):
>>>         return self.dense1(x)
>>> net = Net1In1OutTensor()
>>> x = ops.ones((100, 2), ms.float32)
>>> params = net.trainable_params()
>>> out = net(x)
>>> jw = JacobianWeights(net, out.shape)
>>> jacobian_weights = jw(x, params[0])
>>> print(jacobian_weights.shape)
(100, 1, 1, 2)