assert(llvm::isa<mlir::TFL::ControlNodeOp>(owner));
      assert(control_nodes_at.find(owner) != control_nodes_at.end());
      // Control edge in terms of offsets.
      control_edges.emplace_back(control_nodes_at[outer_op],
                                 control_nodes_at[owner]);
    }
    control_token.dropAllUses();

    // Replace the ControlNodeOp with the wrapped operation.