1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Tensor-like objects that are composed from tf.Tensors.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import abc 22 23import six 24 25from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import 26from tensorflow.python.util import _pywrap_utils 27from tensorflow.python.util import nest 28from tensorflow.python.util.tf_export import tf_export 29 30 31@tf_export("__internal__.CompositeTensor", v1=[]) 32@six.add_metaclass(abc.ABCMeta) 33class CompositeTensor(object): 34 """Abstract base class for Tensor-like objects that are composed from Tensors. 35 36 Each `CompositeTensor` can be decomposed into a structured collection of 37 component `tf.Tensor`s, and reconstructed from those components. 38 39 The `tensorflow.python.util.nest` module has support for treating composite 40 tensors as structure, which makes it easy to flatten and reconstruct 41 composite tensors (or larger structures that contain composite tensors). 42 E.g.: 43 44 ```python 45 ct = ... # Create a composite tensor. 46 flat_list_of_tensors = nest.flatten(ct, expand_composites=True) 47 transformed_list_of_tensors = ... # do something with the flat tensors. 48 result = nest.pack_sequence_as(ct, transformed_list_of_tensors, 49 expand_composites=True) 50 ``` 51 """ 52 53 @abc.abstractproperty 54 def _type_spec(self): 55 """A `TypeSpec` describing the type of this value.""" 56 raise NotImplementedError("%s._type_spec()" % type(self).__name__) 57 58 def _shape_invariant_to_type_spec(self, shape): 59 """Returns a TypeSpec given a shape invariant (used by `tf.while_loop`). 60 61 Args: 62 shape: A `tf.TensorShape` object. The shape invariant for this 63 `CompositeTensor`, or `None` if a default shape invariant should be used 64 (based on the value of this `CompositeTensor`). 65 66 Returns: 67 A nested structure whose values are `tf.TensorShape` objects, specifying 68 the shape invariants for the tensors that comprise this `CompositeTensor`. 69 """ 70 # New TypeSpec subclasses generally do not need to implement this -- 71 # this method is used for backwards compatibility. Users of tf.while_loop 72 # can specify a type by passing in TypeSpec instead. 73 raise NotImplementedError("%s._shape_invariant_to_type_spec" % 74 type(self).__name__) 75 76 def _consumers(self): 77 """Returns a list of `Operation`s that consume this `CompositeTensor`. 78 79 Returns: 80 A list of `Operation`s. 81 82 Raises: 83 RuntimeError: If this method is called while executing eagerly. 84 """ 85 consumers = nest.flatten([ 86 component.consumers() 87 for component in nest.flatten(self, expand_composites=True) 88 if getattr(component, "graph", None) is not None 89 ]) 90 return list(set(consumers)) 91 92 93_pywrap_utils.RegisterType("CompositeTensor", CompositeTensor) 94 95 96def replace_composites_with_components(structure): 97 """Recursively replaces CompositeTensors with their components. 98 99 Args: 100 structure: A `nest`-compatible structure, possibly containing composite 101 tensors. 102 103 Returns: 104 A copy of `structure`, where each composite tensor has been replaced by 105 its components. The result will contain no composite tensors. 106 Note that `nest.flatten(replace_composites_with_components(structure))` 107 returns the same value as `nest.flatten(structure)`. 108 """ 109 if isinstance(structure, CompositeTensor): 110 return replace_composites_with_components( 111 structure._type_spec._to_components(structure)) # pylint: disable=protected-access 112 elif not nest.is_sequence(structure): 113 return structure 114 else: 115 return nest.map_structure( 116 replace_composites_with_components, structure, expand_composites=False) 117 118 119# @TODO(edloper): Can we replace convert_to_tensor_or_xyz with just 120# convert_to_tensor_or_composite? Alternatively, should composite tensors 121# register a dispatch override for tf.convert_to_tensor? 122