Fix control flow bug.

Without this change, the newly added test would fail with a "Reval[0]"
error. This was due to extra Enter nodes being added in the gradients
graph in order to enter the outer while contexts for a value that was
already in those contexts. The extra Enter nodes would cause the
frames of the execution to be messed up, which prevented the final
value from being propagated to the outermost Exit node.

A similar change is probably needed in WhileContext, although I don't
have a test case for this yet.

PiperOrigin-RevId: 186544961