Turns out I was not finished rolling out the Tan improvement.
Have just implemented Tan improvement for models with semi-exogenous shocks. While you cannot apply the Tan improvement to the semi-exogenous shocks themselves, if the model also includes exogenous shocks then you can apply the Tan improvement to these.
A brief explanation: start with standard Tan improvement, we want to go from current agent distribution on (a,z) to next period agent dist on (a’,z’). Before Tan, the standard approach is to use the policy, g, and exogenous shock transition matrix/kernel, pi_z, together to create T which is a markov transition matrix/kernel from (a,z) to (a’,z’). Tan improvement is to do conceptually the same in a computationally more efficient way. First step, create G from g, which goes from (a,z) to (a’,z), and then multiply this by the original distribution. Second step, just use pi_z to go from (a’,z) to (a’,z’).
When we have semi-exogenous shocks we want to get from the current agent distribution on (a,semiz,z) to the next period agent dist on (a’,semiz’,z’). Before Tan, the standard approach is to use the policy, g, and exogenous shock transition matrix/kernel, pi_z, and the semi-exogenous transitions, pi_semiz (which depends on g), together to create T which is a markov transition matrix/kernel from (a,semiz,z) to (a’,semiz’,z’). Tan improvement is to do conceptually the same in a computationally more efficient way. First step, create G from g and pi_semiz, which goes from (a,semiz,z) to (a’,semiz’,z), and then multiply this by the original distribution. Second step, just use pi_z to go from (a’,semiz’,z) to (a’,semiz’,z’). Because semiz is typically only a few grid points (employed/unemployed, or a handful of children), G is still very sparse and so we get a big speed boost (as well as cutting a memory bottleneck).
Note that implementing this Tan improvement for semi-exogenous shocks efficiently requires the ordering (a,semiz,z). Hence I have also had to make a breaking change to adopt this ordering (previously the toolkit was using (a,z,semiz)).