Searched refs:grad_state (Results 1 – 5 of 5) sorted by relevance
211 elif (while_ctxt.grad_state and212 IsContainingContext(while_ctxt.grad_state.forward_context,221 elif (while_ctxt.grad_state and222 while_ctxt.grad_state.forward_context is228 elif (input_while_ctxt.grad_state and229 input_while_ctxt.grad_state.forward_context is while_ctxt):234 elif (input_while_ctxt.grad_state and235 input_ctxt.grad_state.forward_context.grad_state and236 input_ctxt.grad_state.forward_context.grad_state.forward_context is
48 merge_grad = grad_ctxt.grad_state.switch_map.get(op)66 grad_ctxt.grad_state.switch_map[op] = merge_grad105 if grad_ctxt and grad_ctxt.grad_state:110 grad_state = grad_ctxt.grad_state111 real_pred = grad_state.history_map.get(pred.name)114 grad_ctxt = grad_state.grad_context116 history_pred = grad_state.AddForwardAccumulator(pred)121 real_pred = grad_state.AddBackpropAccumulatedValue(history_pred, pred)122 grad_state.history_map[pred.name] = real_pred155 if op._get_control_flow_context().grad_state:[all …]
860 grad_state=self)874 grad_state=self)1096 grad_state = self1098 while pred is None and grad_state:1099 pred = grad_state.history_map.get(cond_ctxt.pred.name)1100 grad_state = grad_state.outer_grad_state1227 for _, grad_state in self._map.items():1229 for y in grad_state.forward_loop_exits:1231 grad_state.pending_exits_count -= 11233 grad_state.unused_exits.append(y)[all …]
679 grad_state = loop_state.GetGradState(x.op, before=False)680 grad_state.deferred_exits.append(x)681 grad_state.pending_exits_count -= 1682 if grad_state.pending_exits_count == 0:685 for y in grad_state.deferred_exits:690 grad_state.unused_exits.append(y)694 for y in grad_state.unused_exits:700 for y in grad_state.unused_exits:
176 def grad_state(self): member in TPUReplicateContext