• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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