Commit 71b06aca authored by Skye Wanderman-Milne's avatar Skye Wanderman-Milne Committed by TensorFlower Gardener
Browse files

Make captured ResourceVariables work with while_v2 + gradients.

This involves a few changes:

* The gradients implementation doesn't error out if it gets a resource tensor y + grad_y (resource tensors previously only worked if grad_y was None). This is necessary since the while body function has matching resource tensor input/outputs, so we end up taking the gradient of the resource output, which has the default grad_y from the outer gradients call.

* We special-case setting the grad_y to None for TensorArrays only. This is currently done with a whitelist. If some TensorArray ops are missing from this list, we can add them later.

* We use the original input names in the body function signature. This isn't necessary but makes debugging easier (otherwise they all show up as "placeholder" and it's hard to tell which input is which).

* Adds new method func_graph.maybe_captured(tensor) to determine if a tensor represents a captured value.

PiperOrigin-RevId: 219298427
parent e21a4cb1
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment