Creates an object to rewrite a parallel-for loop. Args: loop_var: ops.Tensor output of a Placeholder operation. The value should be an int32 scalar representing the loop iteration number. loop_len: A scalar or scalar Tensor representing the number of iterations the l
(self,
loop_var,
loop_len,
pfor_ops,
all_indices=None,
all_indices_partitioned=False,
pfor_config=None)
| 1086 | """ |
| 1087 | |
| 1088 | def __init__(self, |
| 1089 | loop_var, |
| 1090 | loop_len, |
| 1091 | pfor_ops, |
| 1092 | all_indices=None, |
| 1093 | all_indices_partitioned=False, |
| 1094 | pfor_config=None): |
| 1095 | """Creates an object to rewrite a parallel-for loop. |
| 1096 | |
| 1097 | Args: |
| 1098 | loop_var: ops.Tensor output of a Placeholder operation. The value should |
| 1099 | be an int32 scalar representing the loop iteration number. |
| 1100 | loop_len: A scalar or scalar Tensor representing the number of iterations |
| 1101 | the loop is run for. |
| 1102 | pfor_ops: List of all ops inside the loop body. |
| 1103 | all_indices: If not None, an int32 vector with size `loop_len` |
| 1104 | representing the iteration ids that are still active. These values |
| 1105 | should be unique and sorted. However they may not be contiguous. This is |
| 1106 | typically the case when inside a control flow construct which has |
| 1107 | partitioned the indices of the iterations that are being converted. |
| 1108 | all_indices_partitioned: If True, this object is being constructed from a |
| 1109 | control flow construct where not all the pfor iterations are guaranteed |
| 1110 | to be active. |
| 1111 | pfor_config: PForConfig object used while constructing the loop body. |
| 1112 | """ |
| 1113 | assert isinstance(loop_var, ops.Tensor) |
| 1114 | assert loop_var.op.type == "Placeholder" |
| 1115 | self._loop_var = loop_var |
| 1116 | loop_len_value = tensor_util.constant_value(loop_len) |
| 1117 | if loop_len_value is not None: |
| 1118 | loop_len = loop_len_value |
| 1119 | self._loop_len_vector = array_ops.reshape(loop_len, [1]) |
| 1120 | self._all_indices_partitioned = all_indices_partitioned |
| 1121 | if all_indices_partitioned: |
| 1122 | assert all_indices is not None |
| 1123 | self.all_indices = ( |
| 1124 | math_ops.range(loop_len) if all_indices is None else all_indices) |
| 1125 | |
| 1126 | self._conversion_map = object_identity.ObjectIdentityDictionary() |
| 1127 | self._conversion_map[loop_var] = wrap(self.all_indices, True) |
| 1128 | self._pfor_ops = set(pfor_ops) |
| 1129 | self._pfor_op_ids = set([x._id for x in pfor_ops]) |
| 1130 | self._pfor_config = pfor_config |
| 1131 | |
| 1132 | def op_is_inside_loop(self, op): |
| 1133 | """True if op was created inside the pfor loop body.""" |