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

the trace.warmup_posterior suffers from the same issues of being stuck.

If it helps, these are the logps of the initial value (same for both pymc3 and pymc v4)

{'log_k': -0.05, 'log_lambd': -4.52, 'beta': -4.52, 'obs': -4265.52}

I just noticed something very strange - the problem is resolved when I get rid of the dims for "group" - is it possible there’s a bug related to the dims API?

Also note that the point logps of the initial point end up being different with this complete pooling model {'log_k': -0.0, 'log_lambd': -0.23, 'beta': -0.23, 'obs': -4265.52}