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