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

So I checked the shapes and they’re the same between pymc3 and pymc v4 which makes sense because the code is the same. I checked via idata.posterior.dims and also just looking at the shapes of the data variables.

I used the explicit shape argument and it also runs into the same problems as when using dims.

For some reason this model specification just doesnt fit in pymc v4 even though it did in pymc3, however it does fit without any group level indexing in pymc v4.

Not sure how to troubleshoot and identify the problem here because it should work