How about support auto-diff for computing the gradient