def fun(x):
    if (x < 2):
        return lax.mul(2, x)
    elif (x < 5):
        return lax.mul(3, x)
    else:
        return lax.mul(4, x)
def fun(x):
    if (x < 2):
        return lax.mul(2, x)
    elif (x < 5):
        return lax.mul(3, x)
    else:
        return lax.mul(4, x)