Hi @kpei, did you make any progress on this? It recently come to our attention that using pm.dist_math.switch might break the gradient. You can see this post for a case study, and here for some recommendation/workaround.
pm.dist_math.switch