I’m able to reproduce your issue with the chosen sampler (in your larger model with Dirichlet).
It has to do with OpFromGraph gradient failing, which is not a suprise.There were already problems before: Gradient of OpFromGraph fails · Issue #1 · pymc-devs/pytensor · GitHub