def pack_sequence_as(structure, flat_sequence):
"""Like `nest.pack_sequence_as` but also packs other Tensor-like objects.
Args:
structure: The structure to pack into. May contain Tensors,
CompositeTensors, or TensorArrays.
flat_sequence: An iterable containing tensors.
Returns:
A nested structure.
Raises:
AssertionError if `structure` and `flat_sequence` are not compatible.
"""
flat_sequence = list(flat_sequence)
flattened_structure = nest.flatten(structure, expand_composites=True)
if len(flattened_structure) != len(flat_sequence):
raise ValueError("Mismatch in element count")
for i in range(len(flat_sequence)):
if isinstance(flattened_structure[i], tensor_array_ops.TensorArray):
flat_sequence[i] = tensor_array_ops.build_ta_with_new_flow(
old_ta=flattened_structure[i], flow=flat_sequence[i])
return nest.pack_sequence_as(structure,
flat_sequence,
expand_composites=True)
def mark_as_return(self, tensor):
"""Acts like identity but marks the `Tensor` as a return value.
This will possibly return a copy of the `Tensor`. Usage:
```
with AutomaticControlDependencies() as a:
...
t = a.mark_as_return(t)
_ = ...(t...) # i.e. it's safe to use t here
```
Args:
tensor: the `Tensor` to be marked
Returns:
a copy of the `Tensor`.
"""
if isinstance(tensor, ops.IndexedSlices):
values = array_ops.identity(tensor.values)
indices = array_ops.identity(tensor.indices)
self._returned_tensors.add(indices)
self._returned_tensors.add(values)
return ops.IndexedSlices(values,
indices,
dense_shape=tensor.dense_shape)
elif isinstance(tensor, sparse_tensor.SparseTensor):
values = array_ops.identity(tensor.values)
indices = array_ops.identity(tensor.indices)
self._returned_tensors.add(indices)
self._returned_tensors.add(values)
return sparse_tensor.SparseTensor(indices,
values,
dense_shape=tensor.dense_shape)
elif isinstance(tensor, tensor_array_ops.TensorArray):
flow = array_ops.identity(tensor.flow)
self._returned_tensors.add(flow)
return tensor_array_ops.build_ta_with_new_flow(tensor, flow)
# We want to make the return values depend on the stateful operations, but
# we don't want to introduce a cycle, so we make the return value the result
# of a new identity operation that the stateful operations definitely don't
# depend on.
tensor = array_ops.identity(tensor)
self._returned_tensors.add(tensor)
return tensor
def mark_as_return(self, tensor):
"""Acts like identity but marks the `Tensor` as a return value.
This will possibly return a copy of the `Tensor`. Usage:
```
with AutomaticControlDependencies() as a:
...
t = a.mark_as_return(t)
_ = ...(t...) # i.e. it's safe to use t here
```
Args:
tensor: the `Tensor` to be marked
Returns:
a copy of the `Tensor`.
"""
if isinstance(tensor, ops.IndexedSlices):
values = array_ops.identity(tensor.values)
indices = array_ops.identity(tensor.indices)
self._returned_tensors.add(indices)
self._returned_tensors.add(values)
return ops.IndexedSlices(values, indices, dense_shape=tensor.dense_shape)
elif isinstance(tensor, sparse_tensor.SparseTensor):
values = array_ops.identity(tensor.values)
indices = array_ops.identity(tensor.indices)
self._returned_tensors.add(indices)
self._returned_tensors.add(values)
return sparse_tensor.SparseTensor(
indices, values, dense_shape=tensor.dense_shape)
elif isinstance(tensor, tensor_array_ops.TensorArray):
flow = array_ops.identity(tensor.flow)
self._returned_tensors.add(flow)
return tensor_array_ops.build_ta_with_new_flow(tensor, flow)
# We want to make the return values depend on the stateful operations, but
# we don't want to introduce a cycle, so we make the return value the result
# of a new identity operation that the stateful operations definitely don't
# depend on.
tensor = array_ops.identity(tensor)
self._returned_tensors.add(tensor)
return tensor
def pack_sequence_as(structure, flat_sequence):
"""Like `nest.pack_sequence_as` but also packs other Tensor-like objects.
Args:
structure: The structure to pack into. May contain Tensors, IndexedSlices,
TensorArrays or SparseTensors.
flat_sequence: An iterable containing tensors.
Returns:
A nested structure.
Raises:
AssertionError if `structure` and `flat_sequence` are not compatible.
"""
flattened_structure = nest.flatten(structure)
flat_sequence_with_slices_and_tas = []
index = 0
for t in flattened_structure:
if isinstance(t, ops.IndexedSlices):
if t.dense_shape is not None:
flat_sequence_with_slices_and_tas.append(
ops.IndexedSlices(*flat_sequence[index:index + 3]))
index += 3
else:
flat_sequence_with_slices_and_tas.append(
ops.IndexedSlices(*flat_sequence[index:index + 2]))
index += 2
elif isinstance(t, sparse_tensor.SparseTensor):
flat_sequence_with_slices_and_tas.append(
sparse_tensor.SparseTensor(*flat_sequence[index:index + 3]))
index += 3
elif isinstance(t, tensor_array_ops.TensorArray):
flow = flat_sequence[index]
ta = tensor_array_ops.build_ta_with_new_flow(t, flow)
flat_sequence_with_slices_and_tas.append(ta)
index += 1
else:
flat_sequence_with_slices_and_tas.append(flat_sequence[index])
index += 1
assert len(flattened_structure) == len(flat_sequence_with_slices_and_tas)
return nest.pack_sequence_as(structure, flat_sequence_with_slices_and_tas)
def pack_sequence_as(structure, flat_sequence):
"""Like `nest.pack_sequence_as` but also packs other Tensor-like objects.
Args:
structure: The structure to pack into. May contain Tensors, IndexedSlices,
TensorArrays or SparseTensors.
flat_sequence: An iterable containing tensors.
Returns:
A nested structure.
Raises:
AssertionError if `structure` and `flat_sequence` are not compatible.
"""
flattened_structure = nest.flatten(structure)
flat_sequence_with_slices_and_tas = []
index = 0
for t in flattened_structure:
if isinstance(t, ops.IndexedSlices):
if t.dense_shape is not None:
flat_sequence_with_slices_and_tas.append(
ops.IndexedSlices(*flat_sequence[index:index + 3]))
index += 3
else:
flat_sequence_with_slices_and_tas.append(
ops.IndexedSlices(*flat_sequence[index:index + 2]))
index += 2
elif isinstance(t, sparse_tensor.SparseTensor):
flat_sequence_with_slices_and_tas.append(
sparse_tensor.SparseTensor(*flat_sequence[index:index + 3]))
index += 3
elif isinstance(t, tensor_array_ops.TensorArray):
flow = flat_sequence[index]
ta = tensor_array_ops.build_ta_with_new_flow(t, flow)
flat_sequence_with_slices_and_tas.append(ta)
index += 1
else:
flat_sequence_with_slices_and_tas.append(flat_sequence[index])
index += 1
assert len(flattened_structure) == len(flat_sequence_with_slices_and_tas)
return nest.pack_sequence_as(structure, flat_sequence_with_slices_and_tas)
def pack_sequence_as(structure, flat_sequence):
"""Like `nest.pack_sequence_as` but also packs other Tensor-like objects.
Args:
structure: The structure to pack into. May contain Tensors,
CompositeTensors, or TensorArrays.
flat_sequence: An iterable containing tensors.
Returns:
A nested structure.
Raises:
AssertionError if `structure` and `flat_sequence` are not compatible.
"""
flat_sequence = list(flat_sequence)
flattened_structure = nest.flatten(structure, expand_composites=True)
if len(flattened_structure) != len(flat_sequence):
raise ValueError("Mismatch in element count")
for i in range(len(flat_sequence)):
if isinstance(flattened_structure[i], tensor_array_ops.TensorArray):
flat_sequence[i] = tensor_array_ops.build_ta_with_new_flow(
old_ta=flattened_structure[i], flow=flat_sequence[i])
return nest.pack_sequence_as(structure, flat_sequence, expand_composites=True)
def flow_to_tensor_array(flow, ta): # pylint: disable=missing-docstring
return (tensor_array_ops.build_ta_with_new_flow(ta, flow) if isinstance( # pylint: disable=g-long-ternary
ta, tensor_array_ops.TensorArray) else flow)
def flow_to_tensor_array(flow, ta): # pylint: disable=missing-docstring
return (tensor_array_ops.build_ta_with_new_flow(ta, flow) if isinstance( # pylint: disable=g-long-ternary
ta, tensor_array_ops.TensorArray) else flow)
