Simply the logic for bubbling captured tensors when building cond_v2 grad.

The current logic tries to bubble the forward pass tensor to the outermost
graph. That might not always be do-able e.g. when the cond is inside a while loop
it will need to know accumulator logic for while_loop. So instead, the cond_grad
now captures tensors from the forward If op's graph. When the grad If op is
built these tensors will be appropriately captured by the surrounding FuncGraph.

PiperOrigin-RevId: 215993009

0 files changed, 0 insertions, 0 deletions