jacobian

Contents

jacobian#

class pinnx.grad.jacobian(fn, xs, y=None, x=None, mode='backward', vmap=True)[source]#

Compute Jacobian matrix J as J[i, j] = dy_i / dx_j, where i = 0, …, dim_y - 1 and j = 0, …, dim_x - 1.

Parameters:

  • fn (Callable) – Function to compute the gradient.

  • xs (Dict) – Inputs of the function.

  • mode (str) – The mode of the gradient computation. Choose between ‘backward’ and ‘forward’.

  • x (Union[str, Sequence[str], None]) – i`th row. If `i is None, returns the j`th column J[:, `j].

  • y (Union[str, Sequence[str], None]) – j`th column. If `j is None, returns the i`th row J[`i, :], i.e., the gradient of y_i. i and j cannot be both None, unless J has only one element, which is returned.

Returns:

], or j`th column J[:, `j].

Return type:

(i, j)th entry J[i, j], i`th row J[`i,

Contents