Commit 213d76a6 authored by Saurabh Saxena's avatar Saurabh Saxena Committed by TensorFlower Gardener
Browse files

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
parent 4831740f
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment