def fun(x):
    if (x < 3):
        return (x, x)
    else:
        y = lax.mul(2, x)
        return (y, lax.mul(2, y))
def fun(x):
    if (x < 3):
        return (x, x)
    else:
        y = lax.mul(2, x)
        return (y, lax.mul(2, y))