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 // This file defines helper routines for the XLA device. 17 18 #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_ 19 #define TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_ 20 21 #include "absl/types/span.h" 22 #include "tensorflow/compiler/tf2xla/xla_context.h" 23 #include "tensorflow/compiler/xla/client/xla_builder.h" 24 #include "tensorflow/core/framework/tensor.h" 25 26 namespace tensorflow { 27 28 // Helper methods for building XLA computations. 29 class XlaHelpers { 30 public: 31 // Returns a handle representing the zero value of a scalar 32 // element of data_type. 33 static xla::XlaOp Zero(xla::XlaBuilder* b, DataType data_type); 34 35 // Returns a handle representing the one value of a scalar 36 // element of data_type. 37 static xla::XlaOp One(xla::XlaBuilder* b, DataType data_type); 38 39 // Returns a handle representing the given value of an integer scalar 40 // element of data_type. 41 // Note that unlike One and Zero, does not work on boolean types. 42 static xla::XlaOp IntegerLiteral(xla::XlaBuilder* b, DataType data_type, 43 int64 value); 44 45 // Returns a handle representing the given value of a floating-point scalar 46 // element of data_type. 47 static xla::XlaOp FloatLiteral(xla::XlaBuilder* b, DataType data_type, 48 double value); 49 50 // Reshapes literal 'input' to have 'shape'. Both the original shape and 51 // 'shape' must contain the same number of elements. 52 static Status ReshapeLiteral(const xla::Literal& input, 53 absl::Span<const int64> shape, 54 xla::Literal* output); 55 56 // Converts `indices` into a one-hot representation. `depth` is the size 57 // of the new axis to add. `axis` is the position at which to add the new 58 // axis. `indices_shape` is the shape of `indices`. `on_value` and 59 // `off_value` represent the values to use for the on and off positions, 60 // respectively. 61 static Status OneHot(xla::XlaBuilder* builder, int64 depth, int axis, 62 DataType index_type, const TensorShape& indices_shape, 63 const xla::XlaOp& indices, const xla::XlaOp& on_value, 64 const xla::XlaOp& off_value, xla::XlaOp* one_hot); 65 66 // Certain DataTypes should use increased precision DataTypes when performing 67 // reductions. This function remaps a given DataType to a higher precision 68 // DataType if needed. 69 static DataType SumAccumulationType(const DataType& dtype); 70 71 // A helper for creating a ConvertElementType xla op given a DataType rather 72 // than the xla::PrimitiveType. 73 static xla::XlaOp ConvertElementType(const xla::XlaOp& operand, 74 const DataType new_element_type); 75 }; 76 77 } // end namespace tensorflow 78 79 #endif // TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_ 80