xref: /aosp_15_r20/external/tensorflow/tensorflow/python/autograph/lang/special_functions.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2017 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"""Special functions that only make sense for AutoGraph.
16
17These functions are meant to ensure feature parity between Python and AutoGraph,
18so that the exact same code works in both modes. In general, AutoGraph will
19replace these calls.
20"""
21
22from tensorflow.python.autograph.operators import data_structures
23from tensorflow.python.framework import constant_op
24from tensorflow.python.framework import tensor_util
25
26
27def _validate_list_constructor(elements, element_dtype, element_shape):
28  """Validates the inputs of tensor_list."""
29  if element_dtype is not None and element_shape is not None:
30    return
31  if tensor_util.is_tf_type(elements):
32    return
33  if isinstance(elements, (list, tuple)):
34    if elements:
35      return
36    else:
37      raise ValueError(
38          'element_dtype and element_shape are required when elements are'
39          ' empty')
40
41  raise ValueError(
42      'unknown type for elements: {}; only Tensor, list and tuple are'
43      ' allowed'.format(type(elements)))
44
45
46def match_staging_level(value, like_value):
47  """Casts a value to be staged at the same level as another."""
48  if tensor_util.is_tf_type(like_value):
49    return constant_op.constant(value)
50  return value
51
52
53def tensor_list(elements,
54                element_dtype=None,
55                element_shape=None,
56                use_tensor_array=False):
57  """Creates an tensor list and populates it with the given elements.
58
59  This function provides a more uniform access to tensor lists and tensor
60  arrays, and allows optional initialization.
61
62  Note: this function is a simplified wrapper. If you need greater control,
63  it is recommended to use the underlying implementation directly.
64
65  Args:
66    elements: Iterable[tf.Tensor, ...], the elements to initially fill the list
67        with
68    element_dtype: Optional[tf.DType], data type for the elements in the list;
69        required if the list is empty
70    element_shape: Optional[tf.TensorShape], shape for the elements in the list;
71        required if the list is empty
72    use_tensor_array: bool, whether to use the more compatible but restrictive
73        tf.TensorArray implementation
74  Returns:
75    Union[tf.Tensor, tf.TensorArray], the new list.
76  Raises:
77    ValueError: for invalid arguments
78  """
79  _validate_list_constructor(elements, element_dtype, element_shape)
80  if use_tensor_array:
81    return data_structures.tf_tensor_array_new(elements, element_dtype,
82                                               element_shape)
83  else:
84    return data_structures.tf_tensor_list_new(elements, element_dtype,
85                                              element_shape)
86
87
88def stack(list_or_tensor, element_dtype=None, strict=True):
89  """Stacks the input, if it admits the notion of stacking.
90
91  For example, a list of tensors can be stacked into a larger tensor. This
92  function is similar to tf.stack, but it accepts non-lists and lists of
93  non-tensors as arguments. In the latter case, the function does nothing.
94
95  Args:
96    list_or_tensor: Any
97    element_dtype: tf.DType, optional dtypedtype for the elements in the list.
98        Required if the input is stackable, and the list is untyped.
99    strict: bool, if True an error is raised if the input is not stackable.
100        Otherwise the function is a no-op.
101
102  Returns:
103    Any, if the input is stackable, the result will be a tf.Tensor. Otherwise,
104    if strict=False, the result will be list_or_tensor.
105
106  Raises:
107    ValueError: if strict=True and the input is not stackable.
108  """
109  if strict:
110    def raise_error(x):
111      raise ValueError('%s must be stackable when strict=True' % x)
112    original_call = raise_error
113  else:
114    original_call = lambda x: x
115  return data_structures.list_stack(
116      list_or_tensor,
117      data_structures.ListStackOpts(
118          element_dtype=element_dtype, original_call=original_call))
119