| 244 | } |
| 245 | |
| 246 | xla::XlaOp BuildWhile(XlaOpKernelContext* ctx, |
| 247 | const xla::XlaComputation& wrapped_cond, |
| 248 | const xla::XlaComputation& wrapped_body, |
| 249 | const xla::XlaOp& initial_values, |
| 250 | const std::vector<int>& input_mapping, |
| 251 | const std::vector<bool>& compile_time_const_arg_indices, |
| 252 | int num_compile_time_const_args, |
| 253 | bool has_token_input_output) { |
| 254 | xla::XlaOp while_result = |
| 255 | xla::While(wrapped_cond, wrapped_body, initial_values); |
| 256 | std::vector<xla::XlaOp> padded_while_outputs(ctx->num_outputs()); |
| 257 | int while_result_index = 0; |
| 258 | for (int i = 0; i < ctx->num_inputs(); i++) { |
| 259 | if (!compile_time_const_arg_indices[i]) { |
| 260 | padded_while_outputs[input_mapping[while_result_index]] = |
| 261 | xla::GetTupleElement(while_result, while_result_index); |
| 262 | while_result_index++; |
| 263 | } else { |
| 264 | padded_while_outputs[i] = ctx->Input(i); |
| 265 | } |
| 266 | } |
| 267 | // If `body` has a token output, append it to `padded_while_outputs`. |
| 268 | if (has_token_input_output) { |
| 269 | padded_while_outputs.push_back(xla::GetTupleElement( |
| 270 | while_result, ctx->num_inputs() - num_compile_time_const_args)); |
| 271 | } |
| 272 | return xla::Tuple(ctx->builder(), padded_while_outputs); |
| 273 | } |
| 274 | |
| 275 | } // anonymous namespace |
| 276 | |