You don’t want to use vanilla python if/else in PyMC, as only one of the branches will be used from that point on. You need to use pytensor.ifelse.ifelse or pytensor.tensor.switch.
In general if you have a very nested graph you should probably use a Scan for it: scan – Looping in PyTensor — PyTensor dev documentation which is the pytensor way of representing loops.