Compute Jacobian

Compute the Jacobian matrix of a vector-valued function f(x) = [x1*x2, x1+x2^2] at a given input point.

The Jacobian J[i][j] = df_i / dx_j:

• df1/dx1 = x2, df1/dx2 = x1
• df2/dx1 = 1, df2/dx2 = 2*x2

Input: A 1D tensor x of shape (2,).

Output: A 2D tensor of shape (2, 2) — the Jacobian matrix.

API Reference:

• PyTorch: torch.autograd.functional.jacobian(f, x)
• JAX: jax.jacobian(f)(x)