Introduce Operation._add_while_inputs to allow adding inputs to a While op.
This is in preparation for changing while_v2 to rewrite the forward pass to output intermediates needed by the gradient, instead of outputting all intermediates. Since While ops always have the same inputs and output types, we need to be able to add inputs in addition to adding outputs. PiperOrigin-RevId: 223812986
Loading
Please sign in to comment