"""Runs `loop_fn` `iters` times and stacks the outputs.

flat_loop_fn_dtypes = nest.flatten(loop_fn_dtypes)

is_none_list = []

def while_body(i, *ta_list):

"""Body of while loop."""

fn_conv = autograph.tf_convert(loop_fn, autograph_ctx.control_status_ctx())

fn_output = nest.flatten(fn_conv(i))

if len(fn_output) != len(flat_loop_fn_dtypes):

raise ValueError(

f"Number of expected outputs {len(flat_loop_fn_dtypes)}, does not "

f"match the number of actual outputs {len(fn_output)} from loop_fn: "

f"{loop_fn} with output {fn_output}.")

outputs = []

del is_none_list[:]

is_none_list.extend(x is None for x in fn_output)

for out, ta in zip(fn_output, ta_list):

# TODO(agarwal): support returning Operation objects from loop_fn.

if out is not None:

# out may be a ref tensor, wrap it in identity to get a non-ref tensor.

ta = ta.write(i, out)

outputs.append(ta)

return tuple([i + 1] + outputs)

if parallel_iterations is not None:

extra_args = {"parallel_iterations": parallel_iterations}

else:

extra_args = {}

ta_list = while_loop.while_loop(lambda i, *ta: i < iters, while_body, [0] + [

tensor_array_ops.TensorArray(dtype.base_dtype, iters)

for dtype in flat_loop_fn_dtypes

], **extra_args)[1:]

# TODO(rachelim): enable this for sparse tensors

output = [

None if is_none else ta.stack()

for ta, is_none in zip(ta_list, is_none_list)

]

assert len(output) in (0, len(flat_loop_fn_dtypes))

if not output:

# This may happen for the case where iters == 0.

# Pack a list of empty tensors with the proper ranks to match pfor output on 0 iters

loop_var = array_ops.placeholder_with_default(0, shape=[])

try:

loop_fn_out = loop_fn(loop_var)

out_shapes = [

[0] + ops.convert_to_tensor(x).shape

for x in nest.flatten(loop_fn_out)

]

output = [

array_ops.zeros(out_shapes[i], dt)

for i, dt in enumerate(flat_loop_fn_dtypes)

]

except Exception:

output = [array_ops.zeros([0])]

"""Flattens the first two dimensions of x into a single dimension."""

old_shape = array_ops.shape(x)

new_shape = array_ops.concat([[old_shape[0] * old_shape[1]], old_shape[2:]],

axis=0)

"""Check if we are currently inside an XLA compile context."""

g = ops.get_default_graph()

while g is not None:

control_flow_context = g._get_control_flow_context() # pylint: disable=protected-access

while control_flow_context is not None:

if control_flow_context.IsXLAContext():

return True

else:

control_flow_context = control_flow_context.outer_context

# If g is a FuncGraph, get its outer_graph.

g = getattr(g, "outer_graph", None)

iters,

fallback_to_while_loop=True,

parallel_iterations=None,

warn=False):

"""Equivalent to running `loop_fn` `iters` times and stacking the outputs.

def f():

return _pfor_impl(

loop_fn,

iters,

fallback_to_while_loop=fallback_to_while_loop,

parallel_iterations=parallel_iterations,

warn=warn)

# Note that we wrap into a tf.function if in eager execution mode or under

# XLA compilation. The latter is so that we don't compile operations like

# tf.placeholder that are created by the loop body.

functions_run_eagerly = None

if context.executing_eagerly() or _is_under_xla_context():

functions_run_eagerly = def_function.functions_run_eagerly()

if functions_run_eagerly:

logging.warning(

"It looks like tf.function behavior was disabled, perhaps using "

"tf.config.run_functions_eagerly. Vectorization "

"primitives (e.g. tf.vectorized_map) require tf.function to work. "

"These primitives will override the disable.")

def_function.run_functions_eagerly(False)

f = def_function.function(f)

outputs = f()

if functions_run_eagerly is not None:

def_function.run_functions_eagerly(functions_run_eagerly)

return (isinstance(value, composite_tensor.CompositeTensor)

# Leave sparse tensors to be converted by `PFor._convert_sparse`.

and not isinstance(value, sparse_tensor.SparseTensor)

"""Converts a CompositeTensor into a list of stackable tensors."""

if _should_expand_composite(value):

spec = value._type_spec

if not isinstance(spec, type_spec.BatchableTypeSpec):

raise ValueError(f"CompositeTensor instance {value} returned from "

"parallel_for or vectorized_map loop body must provide "

f"a `BatchableTypeSpec` (saw: {spec}).")

if is_batched:

return spec._to_batched_tensor_list(value)

return spec._to_tensor_list(value)

preconverted_value,

batch_size):

"""Converts a list of stacked tensors to a batch CompositeTensor."""

if _should_expand_composite(preconverted_value):

batch_type_spec = preconverted_value._type_spec._batch(batch_size)

return batch_type_spec._from_compatible_tensor_list(stacked_tensors)

"""Test if `loop_fn` has a `pfor_config` argument."""

if tf_inspect.isfunction(loop_fn):

argspec = tf_inspect.getargspec(loop_fn)

return PFOR_CONFIG_ARG in argspec.args

elif isinstance(loop_fn, functools.partial):

fn = loop_fn.func

argspec = tf_inspect.getargspec(fn)

return (PFOR_CONFIG_ARG in argspec.args and

PFOR_CONFIG_ARG not in loop_fn.keywords)

else:

loop_class = tf_decorator.unwrap(loop_fn)[1]

if not hasattr(loop_class, "__call__"):

raise ValueError("`loop_fn` object did not have a __call__ method")

argspec = tf_inspect.getargspec(loop_class.__call__)

iters,

fallback_to_while_loop,

parallel_iterations=None,

pfor_config=None,

warn=False):

"""Implementation of pfor."""

assert not context.executing_eagerly()

loop_fn_has_config = _loop_fn_has_config(loop_fn)

existing_ops = set(ops.get_default_graph().get_operations())

iters_value = tensor_util.constant_value(iters)

# Run the loop body

with ops.name_scope("loop_body"):

loop_var = array_ops.placeholder_with_default(0, shape=[])

if loop_fn_has_config:

if pfor_config is None:

pfor_config = PForConfig()

pfor_config._set_iters(iters) # pylint: disable=protected-access

loop_fn_outputs = loop_fn(loop_var, **{PFOR_CONFIG_ARG: pfor_config})

else:

assert pfor_config is None

f = autograph.tf_convert(loop_fn, autograph_ctx.control_status_ctx())

loop_fn_outputs = f(loop_var)

loop_fn_output_tensors = nest.map_structure(_composite_to_tensors,

loop_fn_outputs)

# Convert outputs to Tensor if needed.

tmp_loop_fn_outputs = []

for loop_fn_output in nest.flatten(loop_fn_output_tensors):

if (loop_fn_output is not None and not isinstance(

loop_fn_output,

(ops.Operation, tensor.Tensor, sparse_tensor.SparseTensor))):

if isinstance(loop_fn_output, indexed_slices.IndexedSlices):

logging.warn("Converting %s to a dense representation may make it slow."

" Alternatively, output the indices and values of the"

" IndexedSlices separately, and handle the vectorized"

" outputs directly." % loop_fn_output)

loop_fn_output = ops.convert_to_tensor(loop_fn_output)

else:

loop_fn_output = ops.convert_to_tensor(loop_fn_output)

tmp_loop_fn_outputs.append(loop_fn_output)

loop_fn_output_tensors = nest.pack_sequence_as(loop_fn_output_tensors,

tmp_loop_fn_outputs)

new_ops = set(ops.get_default_graph().get_operations()) - existing_ops

iters = ops.convert_to_tensor(iters)

if parallel_iterations is not None:

if parallel_iterations < 1:

raise ValueError(

"Argument `parallel_iterations` must be None or a positive integer. "

f"Received: {parallel_iterations}.")

if parallel_iterations == 1:

raise ValueError(

"Found `parallel_iterations == 1`. Use `for_loop` instead.")

if iters_value is not None and iters_value < parallel_iterations:

parallel_iterations = None

if parallel_iterations is None:

with ops.name_scope("pfor"):

converter = PFor(

loop_var,

iters,

new_ops,

fallback_to_while_loop=fallback_to_while_loop,

pfor_config=pfor_config,

warn=warn)

flattened_output_tensors = []

for loop_fn_output in nest.flatten(loop_fn_output_tensors):

output = converter.convert(loop_fn_output)

flattened_output_tensors.append(output)

else:

if pfor_config is not None and pfor_config._has_reductions(): # pylint: disable=protected-access

raise ValueError("Setting `parallel_iterations` currently unsupported if "

"reductions across iterations are performed.")

num_tiled_iterations = iters // parallel_iterations

num_remaining_iterations = iters % parallel_iterations

# TODO(agarwal): Avoid calling loop_fn twice. Generate the loop body inside

# a tf.function and extract the graph from there to vectorize it.

with ops.name_scope("pfor_untiled"):

converter = PFor(loop_var, num_remaining_iterations, new_ops,

fallback_to_while_loop=fallback_to_while_loop,

pfor_config=pfor_config)

remaining_output_tensors = []

flattened_output_tensors = nest.flatten(loop_fn_output_tensors)

for loop_fn_output in flattened_output_tensors:

output = converter.convert(loop_fn_output)

remaining_output_tensors.append(output)

with ops.name_scope("pfor_tiled"):

loop_fn_dtypes = [ops.convert_to_tensor(x).dtype

for x in flattened_output_tensors]

def tiled_loop_body(j):

offset = j * parallel_iterations + num_remaining_iterations

def tiled_loop_fn(i, pfor_config=None):

if loop_fn_has_config:

loop_fn_outputs = loop_fn(i + offset, pfor_config=pfor_config)

else:

loop_fn_outputs = loop_fn(i + offset)

return nest.flatten(

# Stacking across iterations requires explicit Tensors.

nest.map_structure(_composite_to_tensors, loop_fn_outputs))

return _pfor_impl(

tiled_loop_fn,

parallel_iterations,

fallback_to_while_loop=fallback_to_while_loop,

pfor_config=pfor_config)

tiled_output_tensors = for_loop(

tiled_loop_body, loop_fn_dtypes,

num_tiled_iterations, parallel_iterations=1)

tiled_output_tensors = [

_flatten_first_two_dims(y) for y in tiled_output_tensors]

with ops.name_scope("pfor"):

if iters_value is None or iters_value % parallel_iterations:

output_tensors = cond.cond(

math_ops.equal(num_remaining_iterations, 0),

lambda: tiled_output_tensors,

lambda: [array_ops.concat([x, y], axis=0) # pylint: disable=g-long-lambda

for x, y in zip(remaining_output_tensors,

tiled_output_tensors)])

else:

output_tensors = tiled_output_tensors

flattened_output_tensors = nest.flatten(output_tensors)

for output, original_output in zip(flattened_output_tensors,

nest.flatten(loop_fn_output_tensors)):

# Restore any shape information lost from tiling.

# TODO(b/174254748): this may not be correct for stacked `variant`s.

output.set_shape(

tensor_shape.TensorShape([iters_value]).concatenate(

original_output.shape))

return nest.map_structure_up_to(

loop_fn_outputs,

functools.partial(_composite_from_tensors, batch_size=iters_value),

nest.pack_sequence_as(loop_fn_output_tensors,

flattened_output_tensors),

"""Wrapper for gather that implicitly broadcasts unit dimensions."""

static_first_dim = tensor_shape.dimension_value(x.shape[0])

if static_first_dim == 1:

i = 0

elif static_first_dim is None:

i = array_ops.where_v2(array_ops.shape(x)[0] > 1, i, 0)

result = array_ops.gather(x, i)

"""Wrapper for gather that handles CompositeTensors."""

if _should_expand_composite(x):

spec = x._type_spec

gathered_tensors = [_broadcasting_gather(t, i)

for t in spec._to_batched_tensor_list(x)]

return spec._unbatch()._from_compatible_tensor_list(gathered_tensors)

"""Parallel map on the list of tensors unpacked from `elems` on dimension 0.

elems = variable_utils.convert_variables_to_tensors(elems)

elems = nest.map_structure(ops.convert_to_tensor,

elems,

expand_composites=True)

def loop_fn(i):

gathered_elems = nest.map_structure(

lambda x: _gather_from_tensor_or_composite(x, i), elems)

return fn(gathered_elems)

# Extract batch size from the maximum first dimension of any element.

flat_elems = nest.flatten(

nest.map_structure(

functools.partial(_composite_to_tensors,

is_batched=True),

elems))

def _get_shape(x):

if x.shape.rank is None:

return None

return x.shape.as_list()[0]

static_first_dims = [_get_shape(elem) for elem in flat_elems]

if any(s is None for s in static_first_dims):

batch_size = math_ops.reduce_max(

[array_ops.shape(elem)[0] for elem in flat_elems])

else:

batch_size = max(static_first_dims)

return pfor(

loop_fn,

batch_size,

fallback_to_while_loop=fallback_to_while_loop,