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
Loading
Please sign in to comment