keras/src/backend/torch/core.py/switch
def switch(index, branches, *operands):
    index = convert_to_tensor(index, "int32")
    index = torch.clamp(index, 0, len(branches) - 1)
    return branches[index](*operands)