Modelling Repeated Experiments until Success

CustomDist provides gradients as long as you use pytensor operations to implement the generative graph (via the dist argument or the logp function (via the logp function). You don’t need to provide both, and dist is preferred because you get everything – logp and forward sampling. Examples here.

Integrating over the varying lengths should be possible in principle, but in practice working with ragged outputs is hard. That’s assuming the output of the distribution you’re interested in is an entire time series, and not some summary of it. Hard to say without knowing details.

1 Like