Would it be a good idea to add the symbolic representation to docs?
When computing hessian, I get same dimension confusion,
for example,
def pow_adder_reducer(x, y):
  print('x: ', x, 'y: ', y)
  return (2 * x.pow(2) + 3 * y.pow(2)).sum()
inputs = (torch.rand(2), torch.rand(2))
torch.autograd.functional.hessian(pow_adder_reducer, inputs)

PyTorch gives,
x:  tensor([0.4406, 0.9028], grad_fn=<ViewBackward>) y:  tensor([0.8790, 0.7052], grad_fn=<ViewBackward>)
((tensor([[4., 0.],
          [0., 4.]]), tensor([[0., 0.],
          [0., 0.]])), (tensor([[0., 0.],
          [0., 0.]]), tensor([[6., 0.],
          [0., 6.]])))

from sympy I get,

Screenshot (531)1346×539 19.3 KB

is it doing something like,
f1 = 2*(x0**2) + 2*(x1**2)
f2 = 3*(y0**3) + 3*(y1**3)

hessian(f1, (x0, x1))

hessian(f1, (y0, y1))

hessian(f2, (x0, x1))

hessian(f2, (y0, y1))

that results in four 2x2 matrices.