def fun(x):

    def false_fun(x):
        y = (2 * x)
        return (y, (2 * y))
    return lax.cond((x < 3), (lambda x: (x, x)), false_fun, x)
def fun(x):

    def false_fun(x):
        y = (2 * x)
        return (y, (2 * y))
    return lax.cond((x < 3), (lambda x: (x, x)), false_fun, x)