Digging a bit further, it seems that CustomDist initalization wants to call the dist function with the following dummy variables:
<Scalar(int64, shape=())> <Vector(int64, shape=(0,))>
<Scalar(int64, shape=())> is fine as alpha parameter, but <Vector(int64, shape=(0,))> as size makes the graph building process of dist fail on the scan function.
dist(T.scalar(dtype=‘int64’), 0)
Out[7]: Subtensor{i}.0dist(1,T.vector(shape=(0,), dtype=‘int64’))
…
AssertionError: (CumOp{None, add}(Alloc.0), ‘float64’)
Haven’t figured out the solution yet 