Handling failed samples in custom likelihood function

Alternative/supplementary question:

Since something like what I proposed above (throwing away bad samples) might break detailed balance anyway, I figured I’d also try using ADVI or SMC to get posterior estimates. Here I am still having a problems, I suspect related to using ifelse. Here’s the relevant chunk of code. B_hat_solved is the system’s perturbation solution, which will be all nan if no solution exists.

    _, _, B_hat_final, _ = cycle_result
    B_hat_solved = B_hat_final[-1]
    pert_failure_flag = at.or_(at.any(at.isnan(B_hat_solved)), at.any(at.isinf(B_hat_solved)))
    
    log_likelihood = ifelse(pert_failure_flag,
                            at.constant(0.0, dtype='float64'),
                            kalman_likelihood(A, B, B_hat_solved, C, D, observed_data))

My understanding from reading the aesara docs is that if I use ifelse with linker=cvm, it will only evaluate one or the other branch of the conditional. This is important because I need to invert B_hat_solved inside of the likelihood function; if the perturbation failed that whole branch will throw errors.

When I compile this function with aesara.function it works as expected, I can give insane parameter values and I get a zero likelihood without any complaining. As soon as I hand the model to ADVI, however, it immediately fails with FloatingPointError: NaN occurred in optimization. Since I’m also getting Linalg warnings about singular matrices, my suspicion is that both branches of the ifelse are being evaluated, even though I changed the linker in my .aesararc.txt.

Any way to attack this problem via ADVI?