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)