/// <summary>
/// Return the grad state for this op if it's in a forward loop context.
/// </summary>
/// <param name="op"></param>
/// <param name="before"></param>
/// <returns></returns>
public GradLoopState GetGradState(Operation op, bool before)
{
ControlFlowContext forward_ctxt = null;
if (before && util.IsLoopExit(op))
{
forward_ctxt = op._get_control_flow_context();
forward_ctxt = forward_ctxt.outer_context;
if (forward_ctxt != null)
{
forward_ctxt = forward_ctxt.GetWhileContext();
}
}
else
{
forward_ctxt = util.GetWhileContext(op);
}
if (forward_ctxt != null)
{
return(_map.get(forward_ctxt));
}
return(null);
}
