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"""Python wrappers for Datasets and Iterators."""
16from __future__ import absolute_import
17from __future__ import division
18from __future__ import print_function
19
20from tensorflow.python.data.ops import dataset_ops
21from tensorflow.python.data.util import structure
22from tensorflow.python.ops import gen_dataset_ops
23from tensorflow.python.util.tf_export import tf_export
24
25
26@tf_export("data.experimental.get_single_element")
27def get_single_element(dataset):
28  """Returns the single element in `dataset` as a nested structure of tensors.
29
30  This function enables you to use a `tf.data.Dataset` in a stateless
31  "tensor-in tensor-out" expression, without creating an iterator.
32  This can be useful when your preprocessing transformations are expressed
33  as a `Dataset`, and you want to use the transformation at serving time.
34
35  For example:
36
37  ```python
38  def preprocessing_fn(input_str):
39    # ...
40    return image, label
41
42  input_batch = ...  # input batch of BATCH_SIZE elements
43  dataset = (tf.data.Dataset.from_tensor_slices(input_batch)
44             .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
45             .batch(BATCH_SIZE))
46
47  image_batch, label_batch = tf.data.experimental.get_single_element(dataset)
48  ```
49
50  Args:
51    dataset: A `tf.data.Dataset` object containing a single element.
52
53  Returns:
54    A nested structure of `tf.Tensor` objects, corresponding to the single
55    element of `dataset`.
56
57  Raises:
58    TypeError: if `dataset` is not a `tf.data.Dataset` object.
59    InvalidArgumentError (at runtime): if `dataset` does not contain exactly
60      one element.
61  """
62  if not isinstance(dataset, dataset_ops.DatasetV2):
63    raise TypeError("`dataset` must be a `tf.data.Dataset` object.")
64
65  # pylint: disable=protected-access
66  return structure.from_compatible_tensor_list(
67      dataset.element_spec,
68      gen_dataset_ops.dataset_to_single_element(
69          dataset._variant_tensor, **dataset._flat_structure))  # pylint: disable=protected-access
70