/* Copyright 2015 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.
|
==============================================================================*/
|
|
#ifndef TENSORFLOW_CORE_FRAMEWORK_TENSOR_SLICE_H_
|
#define TENSORFLOW_CORE_FRAMEWORK_TENSOR_SLICE_H_
|
|
#include <string>
|
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
|
#include "tensorflow/core/framework/tensor_shape.h"
|
#include "tensorflow/core/framework/tensor_slice.pb.h"
|
#include "tensorflow/core/lib/core/status.h"
|
#include "tensorflow/core/lib/core/stringpiece.h"
|
#include "tensorflow/core/lib/gtl/inlined_vector.h"
|
#include "tensorflow/core/platform/logging.h"
|
|
namespace tensorflow {
|
|
// A tensor slice represents a slice of a given tensor. It is represented by a
|
// list of (start, length) pairs, where the size of the list is the rank of the
|
// tensor.
|
|
class TensorSlice {
|
public:
|
// Construct a tensor slice: you have a number of ways:
|
// -- creating an empty slice
|
// -- from just a dimension (in this case it will create a full slice)
|
// -- from an array of pairs of integers.
|
// -- from a TensorSliceProto protocol buffer
|
// -- from a string format of "start,length:start,length..." where each
|
// "start,length" pair represents the slice on one dimension. We allow a
|
// special "-" that means "everything for this dimension". One such example
|
// is: 0,10:-:14,1:-:-
|
TensorSlice() {}
|
explicit TensorSlice(int dim);
|
explicit TensorSlice(const TensorSliceProto& proto);
|
explicit TensorSlice(std::initializer_list<std::pair<int64, int64>> extents);
|
|
static Status Parse(const string& str, TensorSlice* output);
|
static TensorSlice ParseOrDie(const string& str) {
|
TensorSlice ret;
|
Status s = Parse(str, &ret);
|
if (!s.ok()) {
|
LOG(FATAL) << "Could not parse TensorSlice";
|
}
|
return ret;
|
}
|
|
void Clear();
|
|
// Accessors
|
int dims() const { return starts_.size(); }
|
|
int64 start(int d) const {
|
DCHECK_GE(d, 0);
|
DCHECK_LT(d, dims());
|
return starts_[d];
|
}
|
|
int64 length(int d) const {
|
DCHECK_GE(d, 0);
|
DCHECK_LT(d, dims());
|
return lengths_[d];
|
}
|
|
int64 end(int d) const {
|
DCHECK_GE(d, 0);
|
DCHECK_LT(d, dims());
|
return start(d) + length(d);
|
}
|
|
void set_start(int d, int64 x) {
|
DCHECK_GE(d, 0);
|
DCHECK_LT(d, dims());
|
DCHECK_GE(x, 0);
|
starts_[d] = x;
|
}
|
|
void set_length(int d, int64 x) {
|
DCHECK_GE(d, 0);
|
DCHECK_LT(d, dims());
|
lengths_[d] = x;
|
}
|
|
// If we have a full slice along dimension "d".
|
bool IsFullAt(int d) const {
|
return lengths_[d] == kFullExtent && starts_[d] == 0;
|
}
|
|
// If this is a full slice, i.e. IsFullAt(d) for every d.
|
bool IsFull() const;
|
|
// Set the slice to be a full slice of "dim" dimensions
|
void SetFullSlice(int dim);
|
|
// Extend a slice to "dim" dimensions: all the added dimensions are full.
|
// Requires: dim >= dims().
|
void Extend(int dim);
|
|
// Conversion of a TensorSlice to other formats
|
void AsProto(TensorSliceProto* proto) const;
|
string DebugString() const;
|
|
// Fill *indices and *sizes from *this (so that we can use the slice()
|
// function in eigen tensor). We need a tensor shape in case some of the
|
// slices are full slices.
|
// We allow NDIMS to be greater than dims(), in which case we will pad the
|
// higher dimensions with trivial dimensions.
|
template <int NDIMS>
|
void FillIndicesAndSizes(
|
const TensorShape& shape,
|
Eigen::DSizes<Eigen::DenseIndex, NDIMS>* indices,
|
Eigen::DSizes<Eigen::DenseIndex, NDIMS>* sizes) const;
|
|
// Interaction with other TensorSlices.
|
|
// Compute the intersection with another slice and if "result" is not
|
// nullptr, store the results in *result; returns true if there is any real
|
// intersection.
|
bool Intersect(const TensorSlice& other, TensorSlice* result) const;
|
// A short hand.
|
bool Overlaps(const TensorSlice& other) const {
|
return Intersect(other, nullptr);
|
}
|
|
// Equals iff "*this" and "other" are logically equivalent.
|
bool operator==(const TensorSlice& other) const;
|
bool operator!=(const TensorSlice& other) const { return !(*this == other); }
|
|
// Interaction with TensorShape.
|
|
// Slices a shape and stores the result into *result_shape.
|
// Requires that the shape and *this have the same rank.
|
// For example, given a tensor shape of {3, 4, 5}, and a slice of
|
// 1,2:-:0,2, the result shape is {2, 4, 2}.
|
Status SliceTensorShape(const TensorShape& shape,
|
TensorShape* result_shape) const;
|
|
// Given slice "sub" where "sub" is fully contained in *this,
|
// (meaning that the intersection of "sub" and *this equals "sub"), computes
|
// the "relative" slice of "sub" with respect to *this.
|
//
|
// In other words, if we use A>S to denote slicing a shape S with a slice A,
|
// then the function is computing a slice X such that:
|
// X > (this > S) = sub > S
|
// for any shape S.
|
//
|
// In general, along every dimension, the start of the relative slice is the
|
// start of the "sub" slice minus the start of *this; the length of the
|
// relative slice is the length of the "sub" slice.
|
//
|
// For example, say we have a shape of {3, 4, 5}, "this" is 0,2:-:1,2, and
|
// "sub" is 1,1:2:2,1,2, then the related slice is 1,1:2,2:0,2.
|
//
|
// The caller needs to make sure that "sub" is indeed a sub-slice of *this;
|
// otherwise the result is undefined.
|
void ComputeRelative(const TensorSlice& sub, TensorSlice* relative) const;
|
|
// Updates the slice in such a way that it fully covers "other" slice.
|
// Note, "other" slice should refer to the same tensor shape.
|
// Example:
|
// given a slice [2:4, :, 3:] and "other" slice [:, 1:4, 2:4] the
|
// updated slice would be [:, :, 2:]. Here is why:
|
// dim 0: "2:4" U ":" -> ":"
|
// dim 1: ":" U "1-4" -> ":"
|
// dim 2: "3:" U "2:4" -> "2:"
|
void UpdateToCover(const TensorSlice& other);
|
|
// Returns true if the length field was specified in an Extent.
|
static bool HasExtentLength(const TensorSliceProto::Extent& extent);
|
|
// Returns the value of the length field in an Extent, or -1 if it
|
// is not present.
|
static int64 GetExtentLength(const TensorSliceProto::Extent& extent);
|
|
private:
|
// a length value of kFullExtent (-1) means we have a full slice at this
|
// dimension. It's defined in tensor_slice.cc.
|
static const int64 kFullExtent;
|
|
// TODO(yangke): switch to Eigen once it supports variable size arrays.
|
// A value of
|
gtl::InlinedVector<int64, 4> starts_;
|
gtl::InlinedVector<int64, 4> lengths_;
|
};
|
|
template <int NDIMS>
|
void TensorSlice::FillIndicesAndSizes(
|
const TensorShape& shape, Eigen::DSizes<Eigen::DenseIndex, NDIMS>* indices,
|
Eigen::DSizes<Eigen::DenseIndex, NDIMS>* sizes) const {
|
CHECK_EQ(shape.dims(), dims()) << "Incompatible dimensions between shape "
|
<< "slices: shape = " << shape.DebugString()
|
<< ", slice = " << DebugString();
|
CHECK_GE(NDIMS, dims()) << "Asking for a " << NDIMS << "-dim slice from "
|
<< "a slice of dimension " << dims();
|
for (int d = 0; d < dims(); ++d) {
|
if (IsFullAt(d)) {
|
(*indices)[d] = 0;
|
(*sizes)[d] = shape.dim_size(d);
|
} else {
|
(*indices)[d] = starts_[d];
|
(*sizes)[d] = lengths_[d];
|
}
|
}
|
for (int d = dims(); d < NDIMS; ++d) {
|
(*indices)[d] = 0;
|
(*sizes)[d] = 1;
|
}
|
}
|
|
} // namespace tensorflow
|
|
#endif // TENSORFLOW_CORE_FRAMEWORK_TENSOR_SLICE_H_
|
