From 12bda81b3d998ef554c0e0a7c266bca0b4409f13 Mon Sep 17 00:00:00 2001
From: Thai Nguyen <thaink@google.com>
Date: Thu, 28 May 2020 17:25:45 -0700
Subject: [PATCH] Add Buckettize, SparseCross and BoostedTreesBucketize to flex
 whitelist

PiperOrigin-RevId: 313690852
Change-Id: I9dc21e2439d103ea3783652fb38826e0fd76aaaf
---
 tensorflow/core/kernels/BUILD                          |  8 ++++++++
 tensorflow/core/kernels/boosted_trees/BUILD            |  2 ++
 tensorflow/core/kernels/boosted_trees/quantiles/BUILD  | 10 ++++++++++
 tensorflow/core/platform/BUILD                         |  1 +
 tensorflow/lite/delegates/flex/whitelisted_flex_ops.cc |  5 +++++
 5 files changed, 26 insertions(+)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index c1fc17079c8..47b2fa44f57 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -6617,6 +6617,7 @@ filegroup(
         "batch_matmul_op_impl.h",
         "batch_norm_op.h",
         "broadcast_to_op.h",
+        "bucketize_op.h",
         "control_flow_ops.h",
         "conv_2d.h",
         "conv_3d.h",
@@ -6681,6 +6682,8 @@ filegroup(
         "transpose_op.h",
         "where_op.h",
         "xent_op.h",
+    ] + [
+        "//tensorflow/core/kernels/boosted_trees/quantiles:weighted_quantiles_hdrs",
     ],
 )
 
@@ -6810,6 +6813,7 @@ filegroup(
     srcs = [
         "batchtospace_op.cc",
         "broadcast_to_op.cc",
+        "bucketize_op.cc",
         "ctc_decoder_ops.cc",
         "decode_bmp_op.cc",
         "depthtospace_op.cc",
@@ -6886,6 +6890,7 @@ filegroup(
         "spacetobatch_functor.cc",
         "spacetobatch_op.cc",
         "spacetodepth_op.cc",
+        "sparse_cross_op.cc",
         "sparse_fill_empty_rows_op.cc",
         "sparse_reshape_op.cc",
         "sparse_to_dense_op.cc",
@@ -6934,6 +6939,8 @@ filegroup(
         "where_op.cc",
         "xent_op.cc",
         ":android_extended_ops_headers",
+    ] + [
+        "//tensorflow/core/kernels/boosted_trees:quantile_ops.cc",
     ],
 )
 
@@ -7096,6 +7103,7 @@ cc_library(
     deps = [
         "//tensorflow/core:portable_tensorflow_lib_lite",
         "//tensorflow/core:protos_all_cc_impl",
+        "//tensorflow/core/platform:strong_hash",
         "//third_party/eigen3",
         "//third_party/fft2d:fft2d_headers",
         "@com_google_absl//absl/base",
diff --git a/tensorflow/core/kernels/boosted_trees/BUILD b/tensorflow/core/kernels/boosted_trees/BUILD
index b1f6d9d595f..4ca40dc4177 100644
--- a/tensorflow/core/kernels/boosted_trees/BUILD
+++ b/tensorflow/core/kernels/boosted_trees/BUILD
@@ -19,6 +19,8 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+exports_files(["quantile_ops.cc"])
+
 tf_proto_library(
     name = "boosted_trees_proto",
     srcs = [
diff --git a/tensorflow/core/kernels/boosted_trees/quantiles/BUILD b/tensorflow/core/kernels/boosted_trees/quantiles/BUILD
index 1b48065d6f7..fb03e284d8d 100644
--- a/tensorflow/core/kernels/boosted_trees/quantiles/BUILD
+++ b/tensorflow/core/kernels/boosted_trees/quantiles/BUILD
@@ -10,6 +10,16 @@ package(
 
 exports_files(["LICENSE"])
 
+filegroup(
+    name = "weighted_quantiles_hdrs",
+    srcs = [
+        "quantile_stream_resource.h",
+        "weighted_quantiles_buffer.h",
+        "weighted_quantiles_stream.h",
+        "weighted_quantiles_summary.h",
+    ],
+)
+
 # Quantiles
 
 cc_library(
diff --git a/tensorflow/core/platform/BUILD b/tensorflow/core/platform/BUILD
index 7f7ca0f06cd..30734a840d1 100644
--- a/tensorflow/core/platform/BUILD
+++ b/tensorflow/core/platform/BUILD
@@ -675,6 +675,7 @@ cc_library(
 cc_library(
     name = "strong_hash",
     hdrs = ["strong_hash.h"],
+    visibility = ["//visibility:public"],
     deps = [
         ":platform",
         ":types",
diff --git a/tensorflow/lite/delegates/flex/whitelisted_flex_ops.cc b/tensorflow/lite/delegates/flex/whitelisted_flex_ops.cc
index b38a66f5687..d9150698298 100644
--- a/tensorflow/lite/delegates/flex/whitelisted_flex_ops.cc
+++ b/tensorflow/lite/delegates/flex/whitelisted_flex_ops.cc
@@ -66,6 +66,8 @@ bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
           "BiasAdd",
           "BiasAddGrad",
           "BiasAddV1",
+          "BoostedTreesBucketize",
+          "Bucketize",
           "BroadcastArgs",
           "BroadcastGradientArgs",
           "BroadcastTo",
@@ -386,6 +388,9 @@ bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
           "SparseApplyProximalAdagrad",
           "SparseApplyProximalGradientDescent",
           "SparseApplyRMSProp",
+          "SparseCross",
+          "SparseCrossHashed",
+          "SparseCrossV2",
           "SparseFillEmptyRows",
           "SparseFillEmptyRowsGrad",
           "SparseReshape",