1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15#ifndef TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
16#define TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
17
18#include <memory>
19
20#include "tensorflow/core/framework/tensor_shape.h"
21#include "tensorflow/core/framework/types.pb.h"
22#include "tensorflow/core/platform/refcount.h"
23#include "tensorflow/core/platform/status.h"
24
25namespace tensorflow {
26
27// Abstract interface to a Tensor handle in either tracing or immediate
28// execution mode.
29class AbstractTensorHandle : public core::RefCounted {
30 protected:
31 enum AbstractTensorHandleKind { kGraph, kMlir, kEager, kTfrt, kCustomDevice };
32 explicit AbstractTensorHandle(AbstractTensorHandleKind kind) : kind_(kind) {}
33 ~AbstractTensorHandle() override {}
34
35 public:
36 // Returns tensor dtype.
37 virtual tensorflow::DataType DataType() const = 0;
38
39 // Returns the status of the tensor handle. If it is a tfrt::TensorHandle,
40 // the tensor handle can be an error and return non-OK status.
41 virtual tensorflow::Status TensorHandleStatus() const;
42
43 // Returns tensor shape. If tensor has unknown rank, shape remains untouched.
44 virtual tensorflow::Status Shape(
45 tensorflow::PartialTensorShape* shape) const = 0;
46
47 // The default debug string includes a shape and dtype. Implementations are
48 // free to override it with something more informative.
49 virtual std::string DebugString() const;
50
51 AbstractTensorHandleKind getKind() const { return kind_; }
52
53 private:
54 const AbstractTensorHandleKind kind_;
55};
56
57namespace internal {
58struct AbstractTensorHandleDeleter {
59 void operator()(AbstractTensorHandle* p) const {
60 if (p != nullptr) {
61 p->Unref();
62 }
63 }
64};
65} // namespace internal
66
67// TODO(b/185908092): Make AbstractTensorHandlePtr an IntrusivePtr.
68using AbstractTensorHandlePtr =
69 std::unique_ptr<AbstractTensorHandle,
70 internal::AbstractTensorHandleDeleter>;
71
72} // namespace tensorflow
73
74#endif // TENSORFLOW_C_EAGER_ABSTRACT_TENSOR_HANDLE_H_
75