Update pending count for the inputs of op and enqueue ready ops.
(grads, op, queue, pending_count, loop_state,
xs_set)
| 769 | |
| 770 | |
| 771 | def _UpdatePendingAndEnqueueReady(grads, op, queue, pending_count, loop_state, |
| 772 | xs_set): |
| 773 | """Update pending count for the inputs of op and enqueue ready ops.""" |
| 774 | for x in _NonEagerInputs(op, xs_set): |
| 775 | pending_count[x.op] -= 1 |
| 776 | ready = (pending_count[x.op] == 0) |
| 777 | if loop_state and not ready: |
| 778 | ready = pending_count[x.op] > 0 and control_flow_util.IsLoopSwitch(x.op) |
| 779 | if ready: |
| 780 | if control_flow_util.IsLoopExit(x.op): |
| 781 | # if x is an exit without real gradient, defer processing them. |
| 782 | grad_state = loop_state.GetGradState(x.op, before=False) |
| 783 | grad_state.deferred_exits.append(x) |
| 784 | grad_state.pending_exits_count -= 1 |
| 785 | if grad_state.pending_exits_count == 0: |
| 786 | # We now have all the exits so process them. |
| 787 | has_not_none_grad = False |
| 788 | for y in grad_state.deferred_exits: |
| 789 | if _HasAnyNotNoneGrads(grads, y.op): |
| 790 | has_not_none_grad = True |
| 791 | queue.append(y.op) |
| 792 | else: |
| 793 | grad_state.unused_exits.append(y) |
| 794 | if has_not_none_grad: |
| 795 | # For an unused exit, if it has trainable outputs, backprop |
| 796 | # a zero gradient. Otherwise, just ignore it. |
| 797 | for y in grad_state.unused_exits: |
| 798 | if IsTrainable(y): |
| 799 | _SetGrad(grads, y, loop_state.ZerosLikeForExit(y)) |
| 800 | queue.append(y.op) |
| 801 | else: |
| 802 | # All exits are "unused" so use None as gradient. |
| 803 | for y in grad_state.unused_exits: |
| 804 | queue.append(y.op) |
| 805 | else: |
| 806 | queue.append(x.op) |
| 807 | |
| 808 | |
| 809 | def _SetGrad(grads, t, grad): |
no test coverage detected