You might be trying to force something into a DensityDist that may not fit well the RandomVariable API. I suggest you start with a simple Potential, where you can just define an arbitrary logp expression directly and don’t have to worry about shape inference or nothing of that sort.