def fun(x, y, z):
    pred = lax.lt(x, 3)
    true_fun = (lambda y: y)
    false_fun = (lambda z: lax.neg(z))
    return lax.cond(pred, y, true_fun, z, false_fun)
def fun(x, y, z):
    pred = lax.lt(x, 3)
    true_fun = (lambda y: y)
    false_fun = (lambda z: lax.neg(z))
    return lax.cond(pred, y, true_fun, z, false_fun)