1 /* Copyright 2021 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 // This file declares TF kernel fallback tensor. 17 18 #ifndef TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_TENSOR_H_ 19 #define TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_TENSOR_H_ 20 21 #include "tensorflow/core/framework/tensor.h" 22 #include "tfrt/dtype/dtype.h" // from @tf_runtime 23 #include "tfrt/support/forward_decls.h" // from @tf_runtime 24 #include "tfrt/tensor/tensor.h" // from @tf_runtime 25 #include "tfrt/tensor/tensor_shape.h" // from @tf_runtime 26 27 namespace tensorflow { 28 29 class BaseKernelFallbackTensor : public tfrt::Tensor { 30 public: 31 explicit BaseKernelFallbackTensor(::tensorflow::Tensor tensor); 32 BaseKernelFallbackTensor(const tfrt::TensorShape& shape, tfrt::DType dtype, 33 ::tensorflow::Tensor tensor); 34 35 void Print(tfrt::raw_ostream& os) const override; 36 GetTensor()37 const ::tensorflow::Tensor* GetTensor() const { return &tensor_; } 38 39 private: 40 ::tensorflow::Tensor tensor_; 41 bool is_valid_type_; 42 }; 43 44 class KernelFallbackTensor final 45 : public BaseKernelFallbackTensor, 46 public tfrt::TensorTraits<KernelFallbackTensor> { 47 public: KernelFallbackTensor(::tensorflow::Tensor tensor)48 explicit KernelFallbackTensor(::tensorflow::Tensor tensor) 49 : BaseKernelFallbackTensor(std::move(tensor)) {} KernelFallbackTensor(const tfrt::TensorShape & shape,tfrt::DType dtype,::tensorflow::Tensor tensor)50 KernelFallbackTensor(const tfrt::TensorShape& shape, tfrt::DType dtype, 51 ::tensorflow::Tensor tensor) 52 : BaseKernelFallbackTensor(shape, dtype, std::move(tensor)) {} 53 Create(const tensorflow::Tensor & tensor)54 static KernelFallbackTensor Create(const tensorflow::Tensor& tensor) { 55 return KernelFallbackTensor(tensor); 56 } 57 58 // Tensor type name for KernelFallbackTensor. name()59 static const char* name() { return "KernelFallback"; } 60 }; 61 62 } // namespace tensorflow 63 64 #endif // TENSORFLOW_CORE_RUNTIME_FALLBACK_KERNEL_KERNEL_FALLBACK_TENSOR_H_ 65