def fun(x):
    return lax.while_loop((lambda x: (x < lax.axis_index('i'))), (lambda x: (x + 2)), x)
def fun(x):
    return lax.while_loop((lambda x: (x < lax.axis_index('i'))), (lambda x: (x + 2)), x)