This is why ‘refine’ is not yet a publicly documented feature
Your understanding is exactly right. That is also why this ‘refine’ made solving the Kitao (2008) stationary eqm easy, but did not help for the transition path (in the infinite horizon you can presolve for d, and then use this at every iteration, but for the transition path the d has to be different in every period, so even if you presolve for d you then would only use this once).
Nice that it also speeds up notably for a model similar to Pijoan-Mas (2006). I haven’t actually yet tried using it anywhere except on that Kitao (2008) code so good to know it is something I should maybe think about making the default.
I corrected this by putting,
N_a=prod(n_a);
N_z=prod(n_z);
near start of this function.
VFI Toolkit uses n_a to be a vector of the number of points in each dimension (e.g., if you have two endogenous states with 201 and 101 points respectively, then n_a=[201,101]), whereas when everything then gets vectorized internally N_a is used (so would have N_a=20301, which is 201*101). In your case as there was just one dimension for n_a, you were able to just set N_a=n_a with no issues. Same for n_d vs N_d, n_z vs N_z etc., the lower case is the vector of number of points in each dimension and the upper case is the product of this vector (awkward wording, but it is the number of points in the vectorized problem). All of this is normally hidden from the user, but since you are looking inside you are seeing it all