/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
you may not use this file except in compliance with the License.
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
Unless required by applicable law or agreed to in writing, software
|
distributed under the License is distributed on an "AS IS" BASIS,
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
See the License for the specific language governing permissions and
|
limitations under the License.
|
==============================================================================*/
|
|
#include "tensorflow/core/util/batch_util.h"
|
|
#include "tensorflow/core/framework/register_types.h"
|
#include "tensorflow/core/framework/types.h"
|
#include "tensorflow/core/lib/core/errors.h"
|
|
#define TF_CALL_DATASET_TYPES(m) TF_CALL_ALL_TYPES(m) TF_CALL_QUANTIZED_TYPES(m)
|
|
namespace tensorflow {
|
namespace batch_util {
|
|
namespace {
|
|
Status ValidateInput(const Tensor& parent, const Tensor& element, int64 index) {
|
DCHECK_NE(parent.dim_size(0), 0);
|
DCHECK_GE(index, 0);
|
if (element.NumElements() != (parent.NumElements() / parent.dim_size(0))) {
|
TensorShape chip_shape = parent.shape();
|
chip_shape.RemoveDim(0);
|
return errors::Internal(
|
"ValidateInput Cannot perform copy: number of elements does not match. "
|
" Shapes are: [element]: ",
|
element.shape().DebugString(),
|
", [parent slice]: ", chip_shape.DebugString());
|
}
|
return Status::OK();
|
}
|
|
template <typename T>
|
Status HandleElementToSlice(Tensor element, Tensor* parent, int64 index,
|
bool /* can_move */) {
|
parent->flat_outer_dims<T>().chip(index, 0) = element.flat<T>();
|
return Status::OK();
|
}
|
|
template <>
|
Status HandleElementToSlice<string>(Tensor element, Tensor* parent, int64 index,
|
bool can_move) {
|
auto parent_as_matrix = parent->flat_outer_dims<string>();
|
auto element_flat = element.flat<string>();
|
if (can_move) {
|
for (int64 i = 0; i < element.NumElements(); ++i) {
|
parent_as_matrix(index, i) = std::move(element_flat(i));
|
}
|
} else {
|
parent_as_matrix.chip(index, 0) = element_flat;
|
}
|
return Status::OK();
|
}
|
|
template <>
|
Status HandleElementToSlice<Variant>(Tensor element, Tensor* parent,
|
int64 index, bool can_move) {
|
auto parent_as_matrix = parent->flat_outer_dims<Variant>();
|
auto element_flat = element.flat<Variant>();
|
if (can_move) {
|
for (int64 i = 0; i < element.NumElements(); ++i) {
|
parent_as_matrix(index, i) = std::move(element_flat(i));
|
}
|
} else {
|
parent_as_matrix.chip(index, 0) = element_flat;
|
}
|
return Status::OK();
|
}
|
|
// TODO(b/78245576): Consider removing this overload.
|
template <typename T>
|
void HandleSliceToElement(const Tensor& parent, Tensor* element, int64 index) {
|
element->flat<T>() = parent.flat_outer_dims<T>().chip(index, 0);
|
}
|
|
template <typename T>
|
void HandleSliceToElement(Tensor* parent, Tensor* element, int64 index,
|
bool can_move) {
|
element->flat<T>() = parent->flat_outer_dims<T>().chip(index, 0);
|
}
|
|
template <>
|
void HandleSliceToElement<string>(Tensor* parent, Tensor* element, int64 index,
|
bool can_move) {
|
auto parent_as_matrix = parent->flat_outer_dims<string>();
|
auto element_flat = element->flat<string>();
|
if (can_move) {
|
for (int64 i = 0; i < element->NumElements(); ++i) {
|
element_flat(i) = std::move(parent_as_matrix(index, i));
|
}
|
} else {
|
element_flat = parent_as_matrix.chip(index, 0);
|
}
|
}
|
|
template <>
|
void HandleSliceToElement<Variant>(Tensor* parent, Tensor* element, int64 index,
|
bool can_move) {
|
auto parent_as_matrix = parent->flat_outer_dims<Variant>();
|
auto element_flat = element->flat<Variant>();
|
if (can_move) {
|
for (int64 i = 0; i < element->NumElements(); ++i) {
|
element_flat(i) = std::move(parent_as_matrix(index, i));
|
}
|
} else {
|
element_flat = parent_as_matrix.chip(index, 0);
|
}
|
}
|
|
} // namespace
|
|
// Copies element into the index^th slice of parent (in the 0th dimension).
|
Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index) {
|
TF_RETURN_IF_ERROR(ValidateInput(*parent, element, index));
|
|
bool can_move = element.RefCountIsOne();
|
#define HANDLE_TYPE(T) \
|
case DataTypeToEnum<T>::value: { \
|
return HandleElementToSlice<T>(std::move(element), parent, index, \
|
can_move); \
|
}
|
|
switch (element.dtype()) {
|
TF_CALL_ALL_TYPES(HANDLE_TYPE);
|
TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
|
TF_CALL_uint32(HANDLE_TYPE);
|
TF_CALL_uint64(HANDLE_TYPE);
|
#undef HANDLE_TYPE
|
default:
|
return errors::Unimplemented("CopyElementToSlice Unhandled data type: ",
|
element.dtype());
|
}
|
}
|
|
// Copies the index^th slice of parent (in the 0th dimension) into element.
|
Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index) {
|
TF_RETURN_IF_ERROR(ValidateInput(parent, *element, index));
|
|
#define HANDLE_TYPE(T) \
|
case DataTypeToEnum<T>::value: { \
|
HandleSliceToElement<T>(parent, element, index); \
|
return Status::OK(); \
|
}
|
|
switch (parent.dtype()) {
|
TF_CALL_ALL_TYPES(HANDLE_TYPE);
|
TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
|
#undef HANDLE_TYPE
|
default:
|
return errors::Unimplemented("CopySliceToElement Unhandled data type: ",
|
element->dtype());
|
}
|
}
|
|
// Copies the index^th slice of parent (in the 0th dimension) into element.
|
//
|
// NOTE(mrry): The implementation may be able to optimize the copy to a move.
|
// This is particularly important for DT_STRING tensors.
|
Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element, int64 index) {
|
TF_RETURN_IF_ERROR(ValidateInput(*parent, *element, index));
|
bool can_move = parent->RefCountIsOne();
|
|
#define HANDLE_TYPE(T) \
|
case DataTypeToEnum<T>::value: { \
|
HandleSliceToElement<T>(parent, element, index, can_move); \
|
return Status::OK(); \
|
}
|
|
switch (parent->dtype()) {
|
TF_CALL_ALL_TYPES(HANDLE_TYPE);
|
TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE);
|
#undef HANDLE_TYPE
|
default:
|
return errors::Unimplemented(
|
"MaybeMoveSliceToElement Unhandled data type: ", element->dtype());
|
}
|
}
|
|
// The following five functions are copied from padding_fifo_queue.cc.
|
// TODO(mrry): Reconcile these functions with the similar methods in the
|
// queue implementation.
|
Status ValidateElementToLargerSlice(const Tensor& element, Tensor* parent) {
|
DCHECK_NE(parent->dim_size(0), 0);
|
if (element.NumElements() > (parent->NumElements() / parent->dim_size(0))) {
|
TensorShape chip_shape = parent->shape();
|
chip_shape.RemoveDim(0);
|
return errors::Internal(
|
"HandleElementToLargerSlice Cannot copy slice: number of entries in "
|
"element is greater than number of elements in parent slice. ",
|
"Shapes are: [element]: ", element.shape().DebugString(),
|
", [parent slice]: ", chip_shape.DebugString());
|
}
|
return Status::OK();
|
}
|
|
template <typename T, int NDIMS>
|
Status HandleElementToLargerSlice(const Tensor& element, Tensor* parent,
|
int index) {
|
TF_RETURN_IF_ERROR(ValidateElementToLargerSlice(element, parent));
|
if (element.NumElements() == 0) {
|
return Status::OK();
|
}
|
auto element_t = element.tensor<T, NDIMS>();
|
auto parent_t = parent->tensor<T, NDIMS + 1>();
|
Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_indices;
|
slice_indices[0] = index;
|
Eigen::DSizes<Eigen::DenseIndex, NDIMS + 1> slice_size;
|
slice_size[0] = 1;
|
for (size_t i = 1; i < slice_size.size(); ++i) {
|
slice_size[i] = element_t.dimension(i - 1);
|
}
|
parent_t.slice(slice_indices, slice_size) = element_t.reshape(slice_size);
|
return Status::OK();
|
}
|
|
template <int NDIMS>
|
Status HandleElementToLargerSliceWithRank(const Tensor& element, Tensor* parent,
|
int index) {
|
#define HANDLE_TYPE(T) \
|
case DataTypeToEnum<T>::value: { \
|
return HandleElementToLargerSlice<T, NDIMS>(element, parent, index); \
|
}
|
|
switch (element.dtype()) {
|
TF_CALL_DATASET_TYPES(HANDLE_TYPE);
|
#undef HANDLE_TYPE
|
default:
|
return errors::Unimplemented(
|
"HandleElementToLargerSliceWithRank Unhandled data type: ",
|
element.dtype());
|
}
|
}
|
|
Status CopyElementToLargerSlice(const Tensor& element, Tensor* parent,
|
int index) {
|
if (parent->dims() != element.dims() + 1) {
|
return errors::Internal(
|
"Mismatched ranks. Element's rank is: ", element.dims(),
|
" but element is meant to be a slice in output Tensor having rank: ",
|
parent->dims(), " (should be: ", element.dims() + 1, ")");
|
}
|
|
#define HANDLE_DIMS(NDIMS) \
|
case NDIMS: { \
|
TF_RETURN_IF_ERROR( \
|
HandleElementToLargerSliceWithRank<NDIMS>(element, parent, index)); \
|
return Status::OK(); \
|
}
|
|
switch (element.dims()) {
|
HANDLE_DIMS(0);
|
HANDLE_DIMS(1);
|
HANDLE_DIMS(2);
|
HANDLE_DIMS(3);
|
HANDLE_DIMS(4);
|
HANDLE_DIMS(5);
|
#undef HANDLE_DIMS
|
default:
|
return errors::Unimplemented("CopyElementToLargerSlice Unhandled rank: ",
|
element.dims());
|
}
|
}
|
|
Status SetElementZero(Tensor* element, const Tensor& padding) {
|
#define HANDLE_TYPE(T) \
|
if (element->dtype() == DataTypeToEnum<T>::value) { \
|
element->flat<T>().setConstant(padding.scalar<T>()()); \
|
return Status::OK(); \
|
}
|
TF_CALL_DATASET_TYPES(HANDLE_TYPE);
|
#undef HANDLE_TYPE
|
return errors::Unimplemented("SetElementZero Unhandled data type: ",
|
element->dtype());
|
}
|
|
} // namespace batch_util
|
} // namespace tensorflow
|
