Searched refs:GetGradState (Results 1 – 2 of 2) sorted by relevance
500 def GetGradState(self, op, before): member in _ControlFlowState555 grad_state = self.GetGradState(op, before)561 grad_state = self.GetGradState(op, before)
757 grad_state = loop_state.GetGradState(x.op, before=False)