Time varying survival model (poisson regression) not porting to pymc v4

Questions v5
2

Maybe it is related to the initial value? Could you check in v4 with pm.sample(..., discard_tuned_samples=False) and see if the samples in trace.warmup_posterior also stuck?