Fixes for outside compilation in "If".
1. Add a control edge from predicate sending node to "If" node, so in XlaCompiler we visit "If" node after predicate sending node, and when we visit "If" node, token output for predicate sending node is already available. 2. When lowering "If" node in XLA, the extra token output comes after resource updates, so its index should be "ctx->num_outputs() + num_resource_args" (for "If", we treat each resource arg as a resource update). PiperOrigin-RevId: 225484666
Loading
Please sign in to comment