Incorrect step assignments with custom step function

Questions v5
8

I imagine the sampler fails because of the untransformed Dirichlet? If you’re sampling with NUTS you should definitely use the transform. I removed it to simplify the case with the transformed Dirichlet