def METHOD_NAME(x, cond, local_states):
    # TODO here we could special-case for qubit / ising
    # by taking -x / 1 - x
    # i.e
    # if local_states[0] + local_states[1] == 0:
    #     return jnp.where(cond, -x, x)
    # elif local_states[0] == 0:
    #     return jnp.where(cond, local_states[1] - x, x)
    # elif local_states[1] == 0:
    #     return jnp.where(cond, local_states[0] - x, x)
    # else:
    #     ...
    was_state_0 = x == local_states[0]
    state_0 = jnp.asarray(local_states[0], dtype=x.dtype)
    state_1 = jnp.asarray(local_states[1], dtype=x.dtype)
    return jnp.where(cond ^ was_state_0, state_0, state_1)