Create a PFor object for converting parts of the while_loop. Args: parent_pfor: PFor object being used for converting the while_loop. indices: int32 Tensor of ids for the iterations that are still active (i.e. did not exit the while_loop). cond_stacked: True if the whi
(self, parent_pfor, indices, cond_stacked, inputs,
inputs_stacked)
| 262 | return self |
| 263 | |
| 264 | def _init_pfor(self, parent_pfor, indices, cond_stacked, inputs, |
| 265 | inputs_stacked): |
| 266 | """Create a PFor object for converting parts of the while_loop. |
| 267 | |
| 268 | Args: |
| 269 | parent_pfor: PFor object being used for converting the while_loop. |
| 270 | indices: int32 Tensor of ids for the iterations that are still active |
| 271 | (i.e. did not exit the while_loop). |
| 272 | cond_stacked: True if the while_loop condition is stacked. |
| 273 | inputs: list of input Tensors corresponding 1-to-1 with self._enters. Note |
| 274 | that these Tensors are a subset of the loop variables for the generated |
| 275 | while_loop. |
| 276 | inputs_stacked: List of booleans corresponding 1-to-1 with `inputs`, |
| 277 | indicating if the value is stacked or not. |
| 278 | |
| 279 | Returns: |
| 280 | A PFor instance. The instance is initialized by adding conversion mappings |
| 281 | of nodes that will be external to the conversion that the returned |
| 282 | instance will be used for. e.g. Enter nodes as well as Merge and Switch |
| 283 | outputs are mapped to converted values. |
| 284 | """ |
| 285 | num_outputs = len(self._outputs) |
| 286 | assert len(inputs) == len(self._enters) |
| 287 | assert len(inputs_stacked) == len(self._enters) |
| 288 | loop_var = parent_pfor.loop_var |
| 289 | loop_len = array_ops.size(indices) |
| 290 | pfor = PFor( |
| 291 | loop_var, |
| 292 | loop_len, |
| 293 | pfor_ops=self._pfor_ops, |
| 294 | all_indices=indices, |
| 295 | all_indices_partitioned=cond_stacked, |
| 296 | pfor_config=self._pfor_config) |
| 297 | # Map all inputs of Enter nodes in self._direct_enters to their converted |
| 298 | # values. |
| 299 | for enter in self._direct_enters: |
| 300 | enter_input = enter.op.inputs[0] |
| 301 | converted_enter, stacked, is_sparse_stacked = parent_pfor._convert_helper( |
| 302 | enter_input) |
| 303 | # Since these are resources / variants, they should be unstacked. |
| 304 | assert not stacked and not is_sparse_stacked, (enter, converted_enter) |
| 305 | pfor._add_conversion(enter, wrap(converted_enter, False)) |
| 306 | |
| 307 | # Map all Enter nodes to the inputs. |
| 308 | for enter, inp, stacked in zip(self._enters, inputs, inputs_stacked): |
| 309 | pfor._add_conversion(enter, wrap(inp, stacked)) |
| 310 | # Map outputs of Switch and Merge. |
| 311 | for i in range(num_outputs): |
| 312 | wrapped_inp = wrap(inputs[i], inputs_stacked[i]) |
| 313 | merge = self._enter_merges[i] |
| 314 | pfor._add_conversion(merge.outputs[0], wrapped_inp) |
| 315 | # Note that second output of Merge is typically not used, except possibly |
| 316 | # as a control dependency. To avoid trying to output the correct value, we |
| 317 | # employ a hack here. We output a dummy invalid value with an incorrect |
| 318 | # dtype. This will allow control dependency to work but if using it as an |
| 319 | # input, it should typically lead to errors during graph construction due |
| 320 | # to dtype mismatch. |
| 321 | # TODO(agarwal): Check in the original graph to see if there are any |
no test coverage detected