when I apply Jacobian in PyTorch autograd (torch==1.5.0), it gives me three dimensional output.
for example,
def exp_reducer(x):
    return (x**2).sum(dim=1)
inputs = torch.rand(2, 2)
torch.autograd.functional.jacobian(exp_reducer, inputs)

tensor([[[0.3296, 0.9791],
         [0.0000, 0.0000]],

        [[0.0000, 0.0000],
         [1.6706, 0.7365]]])

and
torch.autograd.functional.jacobian(exp_reducer, inputs).shape

(2,2,2)

but I get a 2 dimensional output when using sympy,
x0, y0, x1, y1= sympy.symbols('x0, y0, x1, y1')
f = Matrix([Function('g')(x0, x1), Function('h')(y0, y1)])
k = Matrix([[x0, x1, y0, y1]])
f.jacobian(k)

then it give me, this,

image1782×293 18.2 KB

should the shape after applying Jacobian be 2d or 3d?