Modifying the code, a key difference will be that do not solve the sparse system of equations for the V_stacked and then reshape to V. Instead, the "solution" of that system should be implicit in the calculation of the residual.
Getting it to auto-differentiate may take some care, as I am not sure if reshape, spdiags, etc. directly support AD without some work. You can use loops instead as required to setup the equations.
Implement the #2 function. This will involve modification of http://www.princeton.edu/~moll/HACTproject/HJB_diffusion_implicit_RBC.m and http://www.princeton.edu/~moll/HACTproject/HJB_no_uncertainty_implicit.m
A few things: