def fun(index, x, y, z):
    branches = [(lambda xyz: xyz[0]), (lambda xyz: lax.neg(xyz[1])), (lambda xyz: lax.sign(xyz[2]))]
    return lax.switch(index, branches, (x, y, z))
def fun(index, x, y, z):
    branches = [(lambda xyz: xyz[0]), (lambda xyz: lax.neg(xyz[1])), (lambda xyz: lax.sign(xyz[2]))]
    return lax.switch(index, branches, (x, y, z))