Implementation of pfor.
(loop_fn, iters, parallel_iterations=None, pfor_config=None)
| 209 | |
| 210 | |
| 211 | def _pfor_impl(loop_fn, iters, parallel_iterations=None, pfor_config=None): |
| 212 | """Implementation of pfor.""" |
| 213 | loop_fn_has_config = _loop_fn_has_config(loop_fn) |
| 214 | existing_ops = set(ops.get_default_graph().get_operations()) |
| 215 | # Run the loop body |
| 216 | with ops.name_scope("loop_body"): |
| 217 | loop_var = array_ops.placeholder(dtypes.int32, shape=[]) |
| 218 | if loop_fn_has_config: |
| 219 | if pfor_config is None: |
| 220 | pfor_config = PForConfig() |
| 221 | pfor_config._set_iters(iters) # pylint: disable=protected-access |
| 222 | loop_fn_outputs = loop_fn(loop_var, **{PFOR_CONFIG_ARG: pfor_config}) |
| 223 | else: |
| 224 | assert pfor_config is None |
| 225 | loop_fn_outputs = loop_fn(loop_var) |
| 226 | |
| 227 | # Convert outputs to Tensor if needed. |
| 228 | tmp_loop_fn_outputs = [] |
| 229 | for loop_fn_output in nest.flatten(loop_fn_outputs): |
| 230 | if (loop_fn_output is not None and not isinstance( |
| 231 | loop_fn_output, |
| 232 | (ops.Operation, ops.Tensor, sparse_tensor.SparseTensor))): |
| 233 | if isinstance(loop_fn_output, indexed_slices.IndexedSlices): |
| 234 | logging.warn("Converting %s to a dense representation may make it slow." |
| 235 | " Alternatively, output the indices and values of the" |
| 236 | " IndexedSlices separately, and handle the vectorized" |
| 237 | " outputs directly." % loop_fn_output) |
| 238 | loop_fn_output = ops.convert_to_tensor(loop_fn_output) |
| 239 | tmp_loop_fn_outputs.append(loop_fn_output) |
| 240 | loop_fn_outputs = nest.pack_sequence_as(loop_fn_outputs, tmp_loop_fn_outputs) |
| 241 | |
| 242 | new_ops = set(ops.get_default_graph().get_operations()) - existing_ops |
| 243 | iters = ops.convert_to_tensor(iters) |
| 244 | if parallel_iterations is not None: |
| 245 | if parallel_iterations < 1: |
| 246 | raise ValueError("parallel_iterations must be None or a positive integer") |
| 247 | if parallel_iterations == 1: |
| 248 | raise ValueError("Found parallel_iterations == 1. Use for_loop instead.") |
| 249 | iters_value = tensor_util.constant_value(iters) |
| 250 | if iters_value is not None and iters_value < parallel_iterations: |
| 251 | parallel_iterations = None |
| 252 | if parallel_iterations is None: |
| 253 | with ops.name_scope("pfor"): |
| 254 | converter = PFor(loop_var, iters, new_ops, pfor_config=pfor_config) |
| 255 | outputs = [] |
| 256 | for loop_fn_output in nest.flatten(loop_fn_outputs): |
| 257 | outputs.append(converter.convert(loop_fn_output)) |
| 258 | return nest.pack_sequence_as(loop_fn_outputs, outputs) |
| 259 | else: |
| 260 | if pfor_config is not None and pfor_config._has_reductions(): # pylint: disable=protected-access |
| 261 | raise ValueError("Setting parallel_iterations currently unsupported if" |
| 262 | " reductions across iterations are performed.") |
| 263 | num_tiled_iterations = iters // parallel_iterations |
| 264 | num_remaining_iterations = iters % parallel_iterations |
| 265 | # TODO(agarwal): Avoid calling loop_fn twice. Generate the loop body inside |
| 266 | # a tf.function and extract the graph from there to vectorize it. |
| 267 | with ops.name_scope("pfor_untiled"): |
| 268 | converter = PFor(loop_var, num_remaining_iterations, new_ops, |
no test coverage detected