1 /* Copyright 2015 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
16 #ifndef TENSORFLOW_CORE_FRAMEWORK_TENSOR_H_
17 #define TENSORFLOW_CORE_FRAMEWORK_TENSOR_H_
18
19 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
20 #include "tensorflow/core/framework/allocator.h"
21 #include "tensorflow/core/framework/tensor_shape.h"
22 #include "tensorflow/core/framework/tensor_types.h"
23 #include "tensorflow/core/framework/types.h"
24 #include "tensorflow/core/framework/types.pb.h"
25 #include "tensorflow/core/lib/core/refcount.h"
26 #include "tensorflow/core/lib/core/status.h"
27 #include "tensorflow/core/lib/core/stringpiece.h"
28 #include "tensorflow/core/lib/gtl/inlined_vector.h"
29 #include "tensorflow/core/platform/logging.h"
30 #include "tensorflow/core/platform/macros.h"
31 #include "tensorflow/core/platform/types.h"
32
33 namespace tensorflow {
34
35 // Forward declarations. In particular, we forward declare protos so that their
36 // symbols can be removed from .so exports.
37 class AllocationDescription;
38 class Allocator;
39 class OpKernelContext;
40 class TensorBuffer;
41 class TensorCApi;
42 class TensorDescription;
43 class TensorProto;
44 class VariantTensorData;
45 namespace batch_util {
46 Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index);
47 } // namespace batch_util
48
49 /// @ingroup core
50 /// Represents an n-dimensional array of values.
51 class Tensor {
52 public:
53 /// \brief Creates a 1-dimensional, 0-element float tensor.
54 ///
55 /// The returned Tensor is not a scalar (shape {}), but is instead
56 /// an empty one-dimensional Tensor (shape {0}, NumElements() ==
57 /// 0). Since it has no elements, it does not need to be assigned a
58 /// value and is initialized by default (IsInitialized() is
59 /// true). If this is undesirable, consider creating a one-element
60 /// scalar which does require initialization:
61 ///
62 /// ```c++
63 ///
64 /// Tensor(DT_FLOAT, TensorShape({}))
65 ///
66 /// ```
67 Tensor();
68
69 /// \brief Creates a Tensor of the given `type` and `shape`. If
70 /// LogMemory::IsEnabled() the allocation is logged as coming from
71 /// an unknown kernel and step. Calling the Tensor constructor
72 /// directly from within an Op is deprecated: use the
73 /// OpKernelConstruction/OpKernelContext allocate_* methods to
74 /// allocate a new tensor, which record the kernel and step.
75 ///
76 /// The underlying buffer is allocated using a `CPUAllocator`.
77 Tensor(DataType type, const TensorShape& shape);
78
79 /// \brief Creates a tensor with the input `type` and `shape`, using
80 /// the allocator `a` to allocate the underlying buffer. If
81 /// LogMemory::IsEnabled() the allocation is logged as coming from
82 /// an unknown kernel and step. Calling the Tensor constructor
83 /// directly from within an Op is deprecated: use the
84 /// OpKernelConstruction/OpKernelContext allocate_* methods to
85 /// allocate a new tensor, which record the kernel and step.
86 ///
87 /// `a` must outlive the lifetime of this Tensor.
88 Tensor(Allocator* a, DataType type, const TensorShape& shape);
89
90 /// \brief Creates a tensor with the input `type` and `shape`, using
91 /// the allocator `a` and the specified "allocation_attr" to
92 /// allocate the underlying buffer. If the kernel and step are known
93 /// allocation_attr.allocation_will_be_logged should be set to true
94 /// and LogMemory::RecordTensorAllocation should be called after the
95 /// tensor is constructed. Calling the Tensor constructor directly
96 /// from within an Op is deprecated: use the
97 /// OpKernelConstruction/OpKernelContext allocate_* methods to
98 /// allocate a new tensor, which record the kernel and step.
99 ///
100 /// `a` must outlive the lifetime of this Tensor.
101 Tensor(Allocator* a, DataType type, const TensorShape& shape,
102 const AllocationAttributes& allocation_attr);
103
104 /// \brief Creates an empty Tensor of the given data type.
105 ///
106 /// Like Tensor(), returns a 1-dimensional, 0-element Tensor with
107 /// IsInitialized() returning True. See the Tensor() documentation
108 /// for details.
109 explicit Tensor(DataType type);
110
111 /// Copy constructor.
112 Tensor(const Tensor& other);
113
114 /// \brief Move constructor. After this call, <other> is safely destructible
115 /// and can be assigned to, but other calls on it (e.g. shape manipulation)
116 /// are not valid.
117 Tensor(Tensor&& other);
118
119 ~Tensor();
120
121 /// Returns the data type.
dtype()122 DataType dtype() const { return shape_.data_type(); }
123
124 /// Returns the shape of the tensor.
shape()125 const TensorShape& shape() const { return shape_; }
126
127 /// \brief Convenience accessor for the tensor shape.
128 ///
129 /// For all shape accessors, see comments for relevant methods of
130 /// `TensorShape` in `tensor_shape.h`.
dims()131 int dims() const { return shape().dims(); }
132
133 /// Convenience accessor for the tensor shape.
dim_size(int d)134 int64 dim_size(int d) const { return shape().dim_size(d); }
135
136 /// Convenience accessor for the tensor shape.
NumElements()137 int64 NumElements() const { return shape().num_elements(); }
138
IsSameSize(const Tensor & b)139 bool IsSameSize(const Tensor& b) const {
140 return shape().IsSameSize(b.shape());
141 }
142
143 // True iff the two tensors use the same underlying refcounted storage
144 bool SharesBufferWith(const Tensor& b) const;
145
146 /// \brief If necessary, has this Tensor been initialized?
147 ///
148 /// Zero-element Tensors are always considered initialized, even if they
149 /// have never been assigned to and do not have any memory allocated.
150 bool IsInitialized() const;
151
152 /// Returns the estimated memory usage of this tensor.
153 size_t TotalBytes() const;
154
155 // Returns the size of sallocated memory for this tensor.
156 size_t AllocatedBytes() const;
157
158 /// Returns true iff this tensor is aligned.
IsAligned()159 bool IsAligned() const {
160 #if EIGEN_MAX_ALIGN_BYTES == 0
161 return true;
162 #else
163 void* ptr = base<void>();
164 return reinterpret_cast<intptr_t>(ptr) % EIGEN_MAX_ALIGN_BYTES == 0;
165 #endif
166 }
167
168 /// Assign operator. This tensor shares other's underlying storage.
169 Tensor& operator=(const Tensor& other) {
170 CopyFromInternal(other, other.shape());
171 return *this;
172 }
173
174 /// Move operator. See move constructor for details.
175 Tensor& operator=(Tensor&& other);
176
177 /// \brief Copy the other tensor into this tensor and reshape it.
178 ///
179 /// This tensor shares other's underlying storage. Returns `true`
180 /// iff `other.shape()` has the same number of elements of the given
181 /// `shape`.
CopyFrom(const Tensor & other,const TensorShape & shape)182 bool CopyFrom(const Tensor& other,
183 const TensorShape& shape) TF_MUST_USE_RESULT {
184 if (other.NumElements() != shape.num_elements()) return false;
185 CopyFromInternal(other, shape);
186 return true;
187 }
188
189 /// \brief Slice this tensor along the 1st dimension.
190
191 /// I.e., the returned tensor satisfies
192 /// returned[i, ...] == this[dim0_start + i, ...].
193 /// The returned tensor shares the underlying tensor buffer with this
194 /// tensor.
195 ///
196 /// NOTE: The returned tensor may not satisfy the same alignment
197 /// requirement as this tensor depending on the shape. The caller
198 /// must check the returned tensor's alignment before calling certain
199 /// methods that have alignment requirement (e.g., `flat()`, `tensor()`).
200 ///
201 /// REQUIRES: `dims()` >= 1
202 /// REQUIRES: `0 <= dim0_start <= dim0_limit <= dim_size(0)`
203 Tensor Slice(int64 dim0_start, int64 dim0_limit) const;
204
205 /// \brief Parse `other` and construct the tensor.
206
207 /// Returns `true` iff the parsing succeeds. If the parsing fails,
208 /// the state of `*this` is unchanged.
209 bool FromProto(const TensorProto& other) TF_MUST_USE_RESULT;
210 bool FromProto(Allocator* a, const TensorProto& other) TF_MUST_USE_RESULT;
211
212 /// \brief Fills in `proto` with `*this` tensor's content.
213 ///
214 /// `AsProtoField()` fills in the repeated field for `proto.dtype()`, while
215 /// `AsProtoTensorContent()` encodes the content in `proto.tensor_content()`
216 /// in a compact form.
217 void AsProtoField(TensorProto* proto) const;
218 void AsProtoTensorContent(TensorProto* proto) const;
219
220 /// \brief Return the tensor data as an `Eigen::Tensor` with the type and
221 /// sizes of this `Tensor`.
222 ///
223 /// Use these methods when you know the data type and the number of
224 /// dimensions of the Tensor and you want an `Eigen::Tensor`
225 /// automatically sized to the `Tensor` sizes. The implementation check
226 /// fails if either type or sizes mismatch.
227 ///
228 /// Example:
229 ///
230 /// ```c++
231 ///
232 /// typedef float T;
233 /// Tensor my_mat(...built with Shape{rows: 3, cols: 5}...);
234 /// auto mat = my_mat.matrix<T>(); // 2D Eigen::Tensor, 3 x 5.
235 /// auto mat = my_mat.tensor<T, 2>(); // 2D Eigen::Tensor, 3 x 5.
236 /// auto vec = my_mat.vec<T>(); // CHECK fails as my_mat is 2D.
237 /// auto vec = my_mat.tensor<T, 3>(); // CHECK fails as my_mat is 2D.
238 /// auto mat = my_mat.matrix<int32>();// CHECK fails as type mismatch.
239 ///
240 /// ```
241 template <typename T>
vec()242 typename TTypes<T>::Vec vec() {
243 return tensor<T, 1>();
244 }
245
246 template <typename T>
matrix()247 typename TTypes<T>::Matrix matrix() {
248 return tensor<T, 2>();
249 }
250
251 template <typename T, size_t NDIMS>
252 typename TTypes<T, NDIMS>::Tensor tensor();
253
254 /// \brief Return the tensor data to an `Eigen::Tensor` with the
255 /// same size but a bitwise cast to the specified dtype `T`.
256 ///
257 /// Using a bitcast is useful for move and copy operations.
258 /// NOTE: this is the same as `tensor()` except a bitcast is allowed.
259 template <typename T, size_t NDIMS>
260 typename TTypes<T, NDIMS>::Tensor bit_casted_tensor();
261
262 /// \brief Return the tensor data to an `Eigen::Tensor` with the
263 /// last dimension elements converted into single elements of a larger type.
264 ///
265 /// For example, this is useful for kernels that can treat NCHW_VECT_C int8
266 /// tensors as NCHW int32 tensors. The sizeof(T) should equal the size of
267 /// the original element type * num elements in the original last dimension.
268 /// NDIMS should be 1 less than the original number of dimensions.
269 template <typename T, size_t NDIMS>
270 typename TTypes<T, NDIMS>::Tensor reinterpret_last_dimension();
271
272 /// \brief Return the tensor data as an `Eigen::Tensor` of the data type and a
273 /// specified shape.
274 ///
275 /// These methods allow you to access the data with the dimensions
276 /// and sizes of your choice. You do not need to know the number of
277 /// dimensions of the Tensor to call them. However, they `CHECK` that
278 /// the type matches and the dimensions requested creates an
279 /// `Eigen::Tensor` with the same number of elements as the tensor.
280 ///
281 /// Example:
282 ///
283 /// ```c++
284 ///
285 /// typedef float T;
286 /// Tensor my_ten(...built with Shape{planes: 4, rows: 3, cols: 5}...);
287 /// // 1D Eigen::Tensor, size 60:
288 /// auto flat = my_ten.flat<T>();
289 /// // 2D Eigen::Tensor 12 x 5:
290 /// auto inner = my_ten.flat_inner_dims<T>();
291 /// // 2D Eigen::Tensor 4 x 15:
292 /// auto outer = my_ten.shaped<T, 2>({4, 15});
293 /// // CHECK fails, bad num elements:
294 /// auto outer = my_ten.shaped<T, 2>({4, 8});
295 /// // 3D Eigen::Tensor 6 x 5 x 2:
296 /// auto weird = my_ten.shaped<T, 3>({6, 5, 2});
297 /// // CHECK fails, type mismatch:
298 /// auto bad = my_ten.flat<int32>();
299 ///
300 /// ```
301 template <typename T>
flat()302 typename TTypes<T>::Flat flat() {
303 return shaped<T, 1>({NumElements()});
304 }
305
306 template <typename T>
unaligned_flat()307 typename TTypes<T>::UnalignedFlat unaligned_flat() {
308 return unaligned_shaped<T, 1>({NumElements()});
309 }
310
311 /// Returns the data as an Eigen::Tensor with NDIMS dimensions, collapsing all
312 /// Tensor dimensions but the last NDIMS-1 into the first dimension of the
313 /// result. If NDIMS > dims() then leading dimensions of size 1 will be
314 /// added to make the output rank NDIMS.
315 template <typename T, size_t NDIMS = 2>
316 typename TTypes<T, NDIMS>::Tensor flat_inner_dims();
317
318 /// Returns the data as an Eigen::Tensor with NDIMS dimensions, collapsing all
319 /// Tensor dimensions but the first NDIMS-1 into the last dimension of the
320 /// result. If NDIMS > dims() then trailing dimensions of size 1 will be
321 /// added to make the output rank NDIMS.
322 template <typename T, size_t NDIMS = 2>
323 typename TTypes<T, NDIMS>::Tensor flat_outer_dims();
324
325 /// Returns the data as an Eigen::Tensor with NDIMS dimensions, collapsing the
326 /// first 'begin' Tensor dimensions into the first dimension of the result and
327 /// the Tensor dimensions of the last dims() - 'begin' - NDIMS into the last
328 /// dimension of the result. If 'begin' < 0 then the |'begin'| leading
329 /// dimensions of size 1 will be added. If 'begin' + NDIMS > dims() then
330 /// 'begin' + NDIMS - dims() trailing dimensions of size 1 will be added.
331 template <typename T, size_t NDIMS = 3>
332 typename TTypes<T, NDIMS>::Tensor flat_inner_outer_dims(int64 begin);
333
334 template <typename T, size_t NDIMS>
335 typename TTypes<T, NDIMS>::Tensor shaped(gtl::ArraySlice<int64> new_sizes);
336
337 /// \brief Return the tensor data to an `Eigen::Tensor` with the new
338 /// shape specified in `new_sizes` and cast to a new dtype `T`.
339 ///
340 /// Using a bitcast is useful for move and copy operations.
341 /// The allowed bitcast is the only difference from `shaped()`.
342 template <typename T, size_t NDIMS>
343 typename TTypes<T, NDIMS>::Tensor bit_casted_shaped(
344 gtl::ArraySlice<int64> new_sizes);
345
346 template <typename T, size_t NDIMS>
347 typename TTypes<T, NDIMS>::UnalignedTensor unaligned_shaped(
348 gtl::ArraySlice<int64> new_sizes);
349
350 /// \brief Return the Tensor data as a `TensorMap` of fixed size 1:
351 /// `TensorMap<TensorFixedSize<T, 1>>`.
352
353 /// Using `scalar()` allows the compiler to perform optimizations as
354 /// the size of the tensor is known at compile time.
355 template <typename T>
356 typename TTypes<T>::Scalar scalar();
357
358 /// Const versions of all the methods above.
359 template <typename T>
vec()360 typename TTypes<T>::ConstVec vec() const {
361 return tensor<T, 1>();
362 }
363
364 template <typename T>
matrix()365 typename TTypes<T>::ConstMatrix matrix() const {
366 return tensor<T, 2>();
367 }
368
369 template <typename T, size_t NDIMS>
370 typename TTypes<T, NDIMS>::ConstTensor tensor() const;
371
372 /// \brief Return the tensor data to an `Eigen::Tensor` with the
373 /// same size but a bitwise cast to the specified dtype `T`.
374 ///
375 /// Using a bitcast is useful for move and copy operations.
376 /// NOTE: this is the same as `tensor()` except a bitcast is allowed.
377 template <typename T, size_t NDIMS>
378 typename TTypes<T, NDIMS>::ConstTensor bit_casted_tensor() const;
379
380 /// \brief Return the tensor data to an `Eigen::Tensor` with the
381 /// last dimension elements converted into single elements of a larger type.
382 ///
383 /// For example, this is useful for kernels that can treat NCHW_VECT_C int8
384 /// tensors as NCHW int32 tensors. The sizeof(T) should equal the size of
385 /// the original element type * num elements in the original last dimension.
386 /// NDIMS should be 1 less than the original number of dimensions.
387 template <typename T, size_t NDIMS>
388 typename TTypes<T, NDIMS>::ConstTensor reinterpret_last_dimension() const;
389
390 template <typename T>
flat()391 typename TTypes<T>::ConstFlat flat() const {
392 return shaped<T, 1>({NumElements()});
393 }
394
395 template <typename T>
unaligned_flat()396 typename TTypes<T>::UnalignedConstFlat unaligned_flat() const {
397 return unaligned_shaped<T, 1>({NumElements()});
398 }
399
400 template <typename T, size_t NDIMS>
401 typename TTypes<T, NDIMS>::ConstTensor shaped(
402 gtl::ArraySlice<int64> new_sizes) const;
403
404 /// \brief Return the tensor data to an `Eigen::Tensor` with the new
405 /// shape specified in `new_sizes` and cast to a new dtype `T`.
406 ///
407 /// Using a bitcast is useful for move and copy operations.
408 /// The allowed bitcast is the only difference from `shaped()`.
409 template <typename T, size_t NDIMS>
410 typename TTypes<T, NDIMS>::ConstTensor bit_casted_shaped(
411 gtl::ArraySlice<int64> new_sizes) const;
412
413 template <typename T, size_t NDIMS>
414 typename TTypes<T, NDIMS>::UnalignedConstTensor unaligned_shaped(
415 gtl::ArraySlice<int64> new_sizes) const;
416
417 template <typename T>
418 typename TTypes<T>::ConstScalar scalar() const;
419
420 template <typename T, size_t NDIMS = 2>
421 typename TTypes<T, NDIMS>::ConstTensor flat_inner_dims() const;
422
423 template <typename T, size_t NDIMS = 2>
424 typename TTypes<T, NDIMS>::ConstTensor flat_outer_dims() const;
425
426 template <typename T, size_t NDIMS = 3>
427 typename TTypes<T, NDIMS>::ConstTensor flat_inner_outer_dims(
428 int64 begin) const;
429
430 /// Render the first `max_entries` values in `*this` into a string.
431 string SummarizeValue(int64 max_entries) const;
432
433 /// A human-readable summary of the tensor suitable for debugging.
434 string DebugString() const;
435
436 /// Fill in the `TensorDescription` proto with metadata about the
437 /// tensor that is useful for monitoring and debugging.
438 void FillDescription(TensorDescription* description) const;
439
440 /// \brief Returns a `StringPiece` mapping the current tensor's buffer.
441 ///
442 /// The returned `StringPiece` may point to memory location on devices
443 /// that the CPU cannot address directly.
444 ///
445 /// NOTE: The underlying tensor buffer is refcounted, so the lifetime
446 /// of the contents mapped by the `StringPiece` matches the lifetime of
447 /// the buffer; callers should arrange to make sure the buffer does
448 /// not get destroyed while the `StringPiece` is still used.
449 ///
450 /// REQUIRES: `DataTypeCanUseMemcpy(dtype())`.
451 StringPiece tensor_data() const;
452
453 /// Copy the other tensor into this tensor and reshape it and reinterpret the
454 /// buffer's datatype.
455 ///
456 /// This tensor shares other's underlying storage.
457 void UnsafeCopyFromInternal(const Tensor&, DataType dtype,
458 const TensorShape&);
459
460 private:
461 // Returns true if the refcount on buf_ and any possible underlying root
462 // buffer is one.
463 bool RefCountIsOne() const;
464 void CheckType(DataType expected_dtype) const;
465 void CheckTypeAndIsAligned(DataType expected_dtype) const;
466 void CheckIsAlignedAndSingleElement() const;
set_dtype(DataType t)467 void set_dtype(DataType t) { shape_.set_data_type(t); }
468
469 // TensorShape's InlineVector.
470 static gtl::InlinedVector<int64, 4> ComputeFlatInnerDims(
471 gtl::ArraySlice<int64> orig, int64 num_out_dims);
472 static gtl::InlinedVector<int64, 4> ComputeFlatOuterDims(
473 gtl::ArraySlice<int64> orig, int64 num_out_dims);
474
475 TensorShape shape_;
476 TensorBuffer* buf_;
477
478 friend class DMAHelper;
479 friend class TensorCApi;
480 friend class TensorReference; // For access to buf_
481 friend class VariableOp; // For access to set_shape
482 friend class AutoReloadVariableOp; // For access to set_shape
483 friend class TensorTestHelper; // For access to set_shape
484 friend class OpKernelContext; // For access to RefCountIsOne().
485 template <typename Device, typename T>
486 friend class AssignVariableOp; // For access to RefCountIsOne().
487 template <typename Device, typename T>
488 friend Status PrepareToUpdateVariable(
489 OpKernelContext* ctx, Tensor* tensor); // For access to RefCountIsOne().
490 friend Status batch_util::CopyElementToSlice(
491 Tensor element, Tensor* parent,
492 int64 index); // For access to RefCountIsOne().
493 friend class NumpyTensorBuffer; // For access to the private constructor
494 // taking the buffer.
495
496 // Creates a tensor with the input datatype, shape and buf.
497 //
498 // Acquires a ref on buf that belongs to this Tensor.
499 Tensor(DataType type, const TensorShape& shape, TensorBuffer* buf);
500
501 bool CanUseDMA() const;
502
503 // Only needed by variable op to set the shape of an uninitialized
504 // Tensor.
505 // TODO: Remove this when we have a better story for detecting
506 // uninitialized tensors.
set_shape(const TensorShape & shape)507 void set_shape(const TensorShape& shape) {
508 DataType dt = dtype();
509 shape_ = shape;
510 set_dtype(dt);
511 }
512
513 void CopyFromInternal(const Tensor& other, const TensorShape& shape);
514
515 template <typename T>
516 T* base() const;
517
518 template <size_t NDIMS>
519 void FillDimsAndValidateCompatibleShape(
520 gtl::ArraySlice<int64> new_sizes,
521 Eigen::array<Eigen::DenseIndex, NDIMS>* dims) const;
522
523 template <typename T, size_t NDIMS>
524 void FillDimsAndValidateCompatibleShape(
525 gtl::ArraySlice<int64> new_sizes,
526 Eigen::array<Eigen::DenseIndex, NDIMS>* dims) const;
527 };
528
529 // Implementation details
530
531 // START_SKIP_DOXYGEN
532
533 // Interface to access the raw ref-counted data buffer.
534 class TensorBuffer : public core::RefCounted {
535 public:
~TensorBuffer()536 ~TensorBuffer() override {}
537
538 // data() points to a memory region of size() bytes.
539 virtual void* data() const = 0;
540 virtual size_t size() const = 0;
541
542 // If this TensorBuffer is sub-buffer of another TensorBuffer,
543 // returns that TensorBuffer. Otherwise, returns this.
544 virtual TensorBuffer* root_buffer() = 0;
545
546 // Fill metadata about the allocation into the proto.
547 virtual void FillAllocationDescription(
548 AllocationDescription* proto) const = 0;
549
550 template <typename T>
base()551 T* base() const {
552 return reinterpret_cast<T*>(data());
553 }
554
555 // Whether this TensorBuffer owns the underlying memory.
OwnsMemory()556 virtual bool OwnsMemory() const { return true; }
557 };
558
559 template <typename T>
base()560 T* Tensor::base() const {
561 return buf_ == nullptr ? nullptr : buf_->base<T>();
562 }
563
564 template <typename T, size_t NDIMS>
tensor()565 typename TTypes<T, NDIMS>::Tensor Tensor::tensor() {
566 CheckTypeAndIsAligned(DataTypeToEnum<T>::v());
567 return typename TTypes<T, NDIMS>::Tensor(base<T>(),
568 shape().AsEigenDSizes<NDIMS>());
569 }
570
571 template <typename T, size_t NDIMS>
tensor()572 typename TTypes<T, NDIMS>::ConstTensor Tensor::tensor() const {
573 CheckTypeAndIsAligned(DataTypeToEnum<T>::v());
574 return typename TTypes<T, NDIMS>::ConstTensor(base<const T>(),
575 shape().AsEigenDSizes<NDIMS>());
576 }
577
578 template <typename T, size_t NDIMS>
bit_casted_tensor()579 typename TTypes<T, NDIMS>::Tensor Tensor::bit_casted_tensor() {
580 CHECK(IsAligned());
581 return typename TTypes<T, NDIMS>::Tensor(base<T>(),
582 shape().AsEigenDSizes<NDIMS>());
583 }
584
585 template <typename T, size_t NDIMS>
bit_casted_tensor()586 typename TTypes<T, NDIMS>::ConstTensor Tensor::bit_casted_tensor() const {
587 CHECK(IsAligned());
588 return typename TTypes<T, NDIMS>::ConstTensor(base<const T>(),
589 shape().AsEigenDSizes<NDIMS>());
590 }
591
592 template <typename T, size_t NDIMS>
reinterpret_last_dimension()593 typename TTypes<T, NDIMS>::Tensor Tensor::reinterpret_last_dimension() {
594 if (NDIMS == dims()) {
595 return tensor<T, NDIMS>();
596 }
597 CHECK(IsAligned());
598 CHECK_EQ(NDIMS, dims() - 1);
599 CHECK_EQ(sizeof(T), shape_.dim_sizes()[NDIMS] * DataTypeSize(dtype()));
600 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
601 for (int d = 0; d < NDIMS; ++d) {
602 dims[d] = shape_.dim_sizes()[d];
603 }
604 return typename TTypes<T, NDIMS>::Tensor(base<T>(), dims);
605 }
606
607 template <typename T, size_t NDIMS>
reinterpret_last_dimension()608 typename TTypes<T, NDIMS>::ConstTensor Tensor::reinterpret_last_dimension()
609 const {
610 if (NDIMS == dims()) {
611 return tensor<T, NDIMS>();
612 }
613 CHECK(IsAligned());
614 CHECK_EQ(NDIMS, dims() - 1);
615 CHECK_EQ(sizeof(T), shape_.dim_sizes()[NDIMS] * DataTypeSize(dtype()));
616 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
617 for (int d = 0; d < NDIMS; ++d) {
618 dims[d] = shape_.dim_sizes()[d];
619 }
620 return typename TTypes<T, NDIMS>::ConstTensor(base<const T>(), dims);
621 }
622
623 template <size_t NDIMS>
FillDimsAndValidateCompatibleShape(gtl::ArraySlice<int64> new_sizes,Eigen::array<Eigen::DenseIndex,NDIMS> * dims)624 void Tensor::FillDimsAndValidateCompatibleShape(
625 gtl::ArraySlice<int64> new_sizes,
626 Eigen::array<Eigen::DenseIndex, NDIMS>* dims) const {
627 CHECK_EQ(NDIMS, new_sizes.size());
628 int64 new_num_elements = 1;
629 for (size_t d = 0; d < NDIMS; d++) {
630 new_num_elements *= new_sizes[d];
631 (*dims)[d] = new_sizes[d];
632 }
633 CHECK_EQ(new_num_elements, NumElements());
634 }
635
636 template <typename T, size_t NDIMS>
FillDimsAndValidateCompatibleShape(gtl::ArraySlice<int64> new_sizes,Eigen::array<Eigen::DenseIndex,NDIMS> * dims)637 void Tensor::FillDimsAndValidateCompatibleShape(
638 gtl::ArraySlice<int64> new_sizes,
639 Eigen::array<Eigen::DenseIndex, NDIMS>* dims) const {
640 CHECK_EQ(NDIMS, new_sizes.size());
641 int64 new_num_elements = 1;
642 for (size_t d = 0; d < NDIMS; d++) {
643 new_num_elements *= new_sizes[d];
644 (*dims)[d] = new_sizes[d];
645 }
646 const int element_size = DataTypeSize(BaseType(dtype()));
647 if (element_size > 0) {
648 CHECK_EQ(new_num_elements * sizeof(T), NumElements() * element_size);
649 } else {
650 // DataTypeSize() returns 0 for some data types. In this case, assume that T
651 // has the same size as the buffer type.
652 // NOTE: If we can be sure that DataTypeSize() does not return 0 for all POD
653 // types, then we should check DataTypeToEnum<T>::v() == dtype(). Or simply
654 // check if `element_size > 0` to err when bit cast is attempted on Tensor
655 // of unknown data type size.
656 CHECK_EQ(new_num_elements, NumElements());
657 }
658 }
659
660 template <typename T, size_t NDIMS>
shaped(gtl::ArraySlice<int64> new_sizes)661 typename TTypes<T, NDIMS>::Tensor Tensor::shaped(
662 gtl::ArraySlice<int64> new_sizes) {
663 CheckTypeAndIsAligned(DataTypeToEnum<T>::v());
664 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
665 FillDimsAndValidateCompatibleShape(new_sizes, &dims);
666 return typename TTypes<T, NDIMS>::Tensor(base<T>(), dims);
667 }
668
669 template <typename T, size_t NDIMS>
bit_casted_shaped(gtl::ArraySlice<int64> new_sizes)670 typename TTypes<T, NDIMS>::Tensor Tensor::bit_casted_shaped(
671 gtl::ArraySlice<int64> new_sizes) {
672 CHECK(IsAligned());
673 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
674 FillDimsAndValidateCompatibleShape<T>(new_sizes, &dims);
675 return typename TTypes<T, NDIMS>::Tensor(base<T>(), dims);
676 }
677
678 template <typename T, size_t NDIMS>
unaligned_shaped(gtl::ArraySlice<int64> new_sizes)679 typename TTypes<T, NDIMS>::UnalignedTensor Tensor::unaligned_shaped(
680 gtl::ArraySlice<int64> new_sizes) {
681 CheckType(DataTypeToEnum<T>::v());
682 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
683 FillDimsAndValidateCompatibleShape(new_sizes, &dims);
684 return typename TTypes<T, NDIMS>::UnalignedTensor(base<T>(), dims);
685 }
686
687 template <typename T, size_t NDIMS>
shaped(gtl::ArraySlice<int64> new_sizes)688 typename TTypes<T, NDIMS>::ConstTensor Tensor::shaped(
689 gtl::ArraySlice<int64> new_sizes) const {
690 CheckType(DataTypeToEnum<T>::v());
691 CHECK(IsAligned());
692 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
693 FillDimsAndValidateCompatibleShape(new_sizes, &dims);
694 return typename TTypes<T, NDIMS>::ConstTensor(base<T>(), dims);
695 }
696
697 template <typename T, size_t NDIMS>
bit_casted_shaped(gtl::ArraySlice<int64> new_sizes)698 typename TTypes<T, NDIMS>::ConstTensor Tensor::bit_casted_shaped(
699 gtl::ArraySlice<int64> new_sizes) const {
700 CHECK(IsAligned());
701 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
702 FillDimsAndValidateCompatibleShape<T>(new_sizes, &dims);
703 return typename TTypes<T, NDIMS>::ConstTensor(base<T>(), dims);
704 }
705
706 template <typename T, size_t NDIMS>
unaligned_shaped(gtl::ArraySlice<int64> new_sizes)707 typename TTypes<T, NDIMS>::UnalignedConstTensor Tensor::unaligned_shaped(
708 gtl::ArraySlice<int64> new_sizes) const {
709 CheckType(DataTypeToEnum<T>::v());
710 Eigen::array<Eigen::DenseIndex, NDIMS> dims;
711 FillDimsAndValidateCompatibleShape(new_sizes, &dims);
712 return typename TTypes<T, NDIMS>::UnalignedConstTensor(base<T>(), dims);
713 }
714
715 template <typename T>
scalar()716 typename TTypes<T>::Scalar Tensor::scalar() {
717 CheckIsAlignedAndSingleElement();
718 return typename TTypes<T>::Scalar(base<T>());
719 }
720
721 template <typename T>
scalar()722 typename TTypes<T>::ConstScalar Tensor::scalar() const {
723 CheckIsAlignedAndSingleElement();
724 return typename TTypes<T>::ConstScalar(base<T>());
725 }
726
727 template <typename T, size_t NDIMS>
flat_inner_dims()728 typename TTypes<T, NDIMS>::Tensor Tensor::flat_inner_dims() {
729 return shaped<T, NDIMS>(ComputeFlatInnerDims(shape_.dim_sizes(), NDIMS));
730 }
731
732 template <typename T, size_t NDIMS>
flat_outer_dims()733 typename TTypes<T, NDIMS>::Tensor Tensor::flat_outer_dims() {
734 return shaped<T, NDIMS>(ComputeFlatOuterDims(shape_.dim_sizes(), NDIMS));
735 }
736
737 template <typename T, size_t NDIMS>
flat_inner_outer_dims(int64 begin)738 typename TTypes<T, NDIMS>::Tensor Tensor::flat_inner_outer_dims(int64 begin) {
739 gtl::InlinedVector<int64, 4> flat_outer =
740 ComputeFlatOuterDims(shape_.dim_sizes(), begin + NDIMS);
741 return shaped<T, NDIMS>(ComputeFlatInnerDims(flat_outer, NDIMS));
742 }
743
744 template <typename T, size_t NDIMS>
flat_inner_dims()745 typename TTypes<T, NDIMS>::ConstTensor Tensor::flat_inner_dims() const {
746 return shaped<T, NDIMS>(ComputeFlatInnerDims(shape_.dim_sizes(), NDIMS));
747 }
748
749 template <typename T, size_t NDIMS>
flat_outer_dims()750 typename TTypes<T, NDIMS>::ConstTensor Tensor::flat_outer_dims() const {
751 return shaped<T, NDIMS>(ComputeFlatOuterDims(shape_.dim_sizes(), NDIMS));
752 }
753
754 template <typename T, size_t NDIMS>
flat_inner_outer_dims(int64 begin)755 typename TTypes<T, NDIMS>::ConstTensor Tensor::flat_inner_outer_dims(
756 int64 begin) const {
757 gtl::InlinedVector<int64, 4> flat_outer =
758 ComputeFlatOuterDims(shape_.dim_sizes(), begin + NDIMS);
759 return shaped<T, NDIMS>(ComputeFlatInnerDims(flat_outer, NDIMS));
760 }
761
Tensor(const Tensor & other)762 inline Tensor::Tensor(const Tensor& other)
763 : shape_(other.shape()), buf_(other.buf_) {
764 if (buf_) buf_->Ref();
765 }
766
Tensor(Tensor && other)767 inline Tensor::Tensor(Tensor&& other)
768 : shape_(std::move(other.shape())), buf_(other.buf_) {
769 other.buf_ = nullptr;
770 }
771
772 inline Tensor& Tensor::operator=(Tensor&& other) {
773 // Avoid self-assignment, since we might destroy our underlying buffer.
774 if (&other != this) {
775 shape_ = std::move(other.shape_);
776 if (buf_) buf_->Unref();
777 buf_ = other.buf_;
778 other.buf_ = nullptr;
779 }
780 return *this;
781 }
782
783 // END_SKIP_DOXYGEN
784
785 } // namespace tensorflow
786
787 #endif // TENSORFLOW_CORE_FRAMEWORK_TENSOR_H_
788