def cond_via_switch(pred, true_fun, false_fun, op, *args):
    if (len(args) > 0):
        assert (len(args) == 1)
        (true_op, _true_fun, false_op, _false_fun) = (true_fun, false_fun, op, args[0])
        op = (false_op, true_op)
        false_fun = (lambda op: _false_fun(op[0]))
        true_fun = (lambda op: _true_fun(op[1]))
    index = lax.convert_element_type(pred, np.int32)
    return lax.switch(index, [false_fun, true_fun], op)
def cond_via_switch(pred, true_fun, false_fun, op, *args):
    if (len(args) > 0):
        assert (len(args) == 1)
        (true_op, _true_fun, false_op, _false_fun) = (true_fun, false_fun, op, args[0])
        op = (false_op, true_op)
        false_fun = (lambda op: _false_fun(op[0]))
        true_fun = (lambda op: _true_fun(op[1]))
    index = lax.convert_element_type(pred, np.int32)
    return lax.switch(index, [false_fun, true_fun], op)