Support nested while_v2's (see change to _get_body_graph).
Fix gradient logic when loop body has an Identity node whose output value is needed for computing a gradient. Since we do not accumulate outputs of Identity nodes we instead try to use the output of the first non-Identity node in the parent hierarchy. Raise exception when taking gradient of a while loop which creates a resource in its body, i.e., the resource is not captured. PiperOrigin-RevId: 218884162
Loading
Please sign in to comment