PyMC3 elementwise if condition

Try aesara.tensor.switch