CustomDist is called twice, once without size, and once with, just to figure out the support dimensionality of the distribution. The second time it’s called with the actual size. So you can try to handle the case where size.type.shape==(0,) just so the second call happens.
There are however two greater issues, in that PyMC will not be able to guess the logp of your dist. The first is that your CustomDist involves two RVs, the first Uniform first_roll and the Categorical choice inside the Scan. Actually I don’t know if this is a problem in general, but definitely the graph between first_roll and assignments is too much for PyMC to figure out how to get the logprob. Either way, you could get around this by moving the first Uniform outside of the CustomDist and passing it as a parameter.
The more important issue is that PyMC does not know how to derive the logprob of a set_subtensor operation. If you can rewrite it to avoid it, it should be able to derive the whole logprob. Also PyMC doesn’t know what to do with the last astype(int) but looks like you should be able to avoid it.
I suggest you start with a simpler Scan RV (even if it’s not what you care about), to put a finger on where things fail. PyMC logprob inference is still new, and has a lot of operations it can’t handle. Unfortunately we don’t have useful messages besides “logprob could not be derived” which are not useful when you have so many operations that could have caused it to fail.