1 /* Copyright 2017 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 // Utilities for working with XLA Literals. 17 18 #ifndef TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 19 #define TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 20 21 #include "absl/types/span.h" 22 #include "tensorflow/compiler/xla/literal.h" 23 #include "tensorflow/compiler/xla/xla_data.pb.h" 24 #include "tensorflow/core/framework/tensor.h" 25 #include "tensorflow/core/lib/core/status.h" 26 27 namespace tensorflow { 28 29 // Returns a BorrowingLiteral that utilizes the same underlying buffer owned by 30 // 'host_tensor'. 31 Status HostTensorToBorrowingLiteral(const Tensor& host_tensor, 32 xla::BorrowingLiteral* literal); 33 34 // Returns a Literal with the contents of 'host_tensor', backed by its own 35 // storage (i.e., not reusing 'host_tensor's buffers.) 36 xla::StatusOr<xla::Literal> HostTensorToLiteral(const Tensor& host_tensor); 37 38 // Returns a MutableBorrowingLiteral that utilizes the same underlying buffer 39 // owned by 'host_tensor', but is mutable via the xla::Literal methods. 40 Status HostTensorToMutableBorrowingLiteral( 41 Tensor* host_tensor, xla::MutableBorrowingLiteral* literal); 42 // Similar as above, except the literal shape is explicitly provided and used 43 // instead of obtaining it from the 'host_tensor'. The provided literal shape 44 // 'xla_shape' must be compatible with the shape of 'host_tensor'. 45 Status HostTensorToMutableBorrowingLiteral( 46 const xla::Shape& xla_shape, Tensor* host_tensor, 47 xla::MutableBorrowingLiteral* literal); 48 49 // Returns a BorrowingLiteral tuple that utilizes the same underlying buffers 50 // owned by 'host_tensors'. 51 Status HostTensorsToBorrowingLiteralTuple(absl::Span<const Tensor> host_tensors, 52 xla::BorrowingLiteral* literal); 53 54 // Copies 'literal' to freshly allocated 'host_tensor', which is allocated of 55 // type <target_type>. 56 // Fails if the literal's primitive type != 57 // DataTypeToPrimitiveType(target_type). Note that <target_type> is not 58 // derivable from the type of <literal>, because multiple tensorflow types map 59 // to the same XLA type (e.g. INT32 and QINT32 both map to INT32 in 60 // XLA). 61 Status LiteralToHostTensor(const xla::LiteralSlice& literal, 62 DataType target_type, Tensor* host_tensor); 63 64 // Copies the contents of 'literal' to a previously allocated tensor 65 // 'host_tensor'. The tensor and the literal must have the same number of 66 // elements and the same type. 67 Status CopyLiteralToHostTensor(const xla::LiteralSlice& literal, 68 Tensor* host_tensor); 69 70 } // namespace tensorflow 71 72 #endif // TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ 73