Method StateAlongEdge

tensorflow/compiler/tf2xla/functionalize_cond.cc:1103–1131

}

StateMap::CondId FunctionalizeCond::StateAlongEdge(const Edge* e) {
  Node* src = e->src();
  StateMap::CondId id = state_map_.LookupCondId(e->src());

  // Dead nodes only propagate dead state.
  if (state_map_.IsDead(id)) return id;

  if (IsSwitch(src)) {
    StateMap::CondState state;
    if (id != nullptr) state = *id;
    OutputTensor predicate;
    TF_CHECK_OK(GetSwitchPredicate(*src, &predicate));
    if (e->IsControlEdge()) {
      // In gradients of tf.cond(), in each branch, we have a NoOp node as
      // control pivot. These NoOp nodes have control dependency from Switch
      // node. If we don't record this into CondState, branches might have
      // incorrect CondState (e.g. if the branch only has a Const data node).
      // We set it to kNeither because there is no way to tell whether it's
      // for true branch or false branch. This node's desendents might have
      // other incoming edges with defined BranchType, and we correctly handle
      // merging kNeither with other defined BranchType in StateAlongEdge().
      state[predicate] = BranchType::kNeither;
    } else {
      state[predicate] = BranchType(e->src_output());
    }
    return state_map_.GetCondId(state);
  }
  return id;
}

Status FunctionalizeCond::DetermineCondStateMerge(Node* dst) {
  // Only Merge nodes with two inputs are supported, but if this is a redundant