1 /* Copyright 2018 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 // XLA-specific ListDiff Op. This only supports constant DT_INT32 and DT_INT64 17 // input. 18 19 #include <unordered_set> 20 21 #include "tensorflow/compiler/tf2xla/type_util.h" 22 #include "tensorflow/compiler/tf2xla/xla_helpers.h" 23 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" 24 #include "tensorflow/compiler/tf2xla/xla_op_registry.h" 25 #include "tensorflow/compiler/xla/client/xla_builder.h" 26 #include "tensorflow/core/framework/kernel_def_builder.h" 27 #include "tensorflow/core/framework/register_types.h" 28 #include "tensorflow/core/lib/core/errors.h" 29 30 namespace tensorflow { 31 namespace { 32 33 constexpr std::array<DataType, 2> kListDiffTypes = {DT_INT32, DT_INT64}; 34 35 // ListDiffOp is an XLA kernel that supports constant-only x and y input. 36 class ListDiffOp : public XlaOpKernel { 37 public: ListDiffOp(OpKernelConstruction * context)38 explicit ListDiffOp(OpKernelConstruction* context) : XlaOpKernel(context) {} 39 Compile(XlaOpKernelContext * context)40 void Compile(XlaOpKernelContext* context) override { 41 OP_REQUIRES(context, TensorShapeUtils::IsVector(context->InputShape(0)), 42 errors::InvalidArgument("ListDiff expects x as a vector, not ", 43 context->InputShape(0).DebugString())); 44 45 OP_REQUIRES(context, TensorShapeUtils::IsVector(context->InputShape(1)), 46 errors::InvalidArgument("ListDiff expects y as a vector, not ", 47 context->InputShape(1).DebugString())); 48 49 DataType val_type = context->expected_output_dtype(0); 50 DataType idx_type = context->expected_output_dtype(1); 51 52 Status status; 53 switch (val_type) { 54 case DT_INT32: 55 status = ListDiffWithIndexType<int32>(context, idx_type); 56 break; 57 case DT_INT64: 58 status = ListDiffWithIndexType<int64_t>(context, idx_type); 59 break; 60 default: 61 // This should never happen since we restrict this kernel to only match 62 // inputs with supported Tensor datatype. 63 status = errors::InvalidArgument("ListDiff expects x and y as either ", 64 "int32 or int64, not ", 65 DataTypeString(val_type)); 66 } 67 OP_REQUIRES_OK(context, status); 68 } 69 70 private: 71 template <typename Tval, typename Tidx> ListDiff(XlaOpKernelContext * context)72 Status ListDiff(XlaOpKernelContext* context) { 73 std::vector<int64_t> x_input, y_input; 74 TF_RETURN_IF_ERROR(context->ConstantInputAsIntVector(0, &x_input)); 75 TF_RETURN_IF_ERROR(context->ConstantInputAsIntVector(1, &y_input)); 76 77 std::unordered_set<Tval> y_input_set; 78 y_input_set.reserve(y_input.size()); 79 for (auto y : y_input) { 80 y_input_set.insert(y); 81 } 82 83 std::vector<Tval> val_output; 84 std::vector<Tidx> idx_output; 85 auto x_size = x_input.size(); 86 for (Tidx i = 0; i < x_size; ++i) { 87 if (y_input_set.count(x_input[i]) > 0) { 88 continue; 89 } 90 val_output.push_back(x_input[i]); 91 idx_output.push_back(i); 92 } 93 94 context->SetOutput(0, 95 xla::ConstantR1<Tval>(context->builder(), val_output)); 96 context->SetOutput(1, 97 xla::ConstantR1<Tidx>(context->builder(), idx_output)); 98 return OkStatus(); 99 } 100 101 template <typename Tval> ListDiffWithIndexType(XlaOpKernelContext * context,DataType idx_type)102 Status ListDiffWithIndexType(XlaOpKernelContext* context, DataType idx_type) { 103 switch (idx_type) { 104 case DT_INT32: 105 return ListDiff<Tval, int32>(context); 106 case DT_INT64: 107 return ListDiff<Tval, int64_t>(context); 108 default: 109 return errors::InvalidArgument( 110 "ListDiff expects idx_out as either int32 or int64, not ", 111 DataTypeString(idx_type)); 112 } 113 } 114 }; 115 116 REGISTER_XLA_OP(Name("ListDiff") 117 .TypeConstraint("T", kListDiffTypes) 118 .CompileTimeConstantInput("x") 119 .CompileTimeConstantInput("y"), 120 ListDiffOp); 121 122 } // namespace 123 } // namespace tensorflow 124