You are probably thinking of forward-mode autodiff (which is more useful when the dimensionality of the output of your function is comparatively large), which i

You are probably thinking of forward-mode autodiff (which is more useful when the dimensionality of the output of your function is comparatively large), which is different from backward-mode autodiff (which is more useful when the dimensionality of the output is comparatively small). Both will work, but one will be more efficient than the other (depending on the context). For things like "training neural networks", people tend to use backward-mode (since you often want to optimize a single loss output w.r.t. many things).