Move flex ops whitelist to a library outside toco
This enables its reuse from other places outside toco. PiperOrigin-RevId: 249776820
This commit is contained in:
parent
de43b973f8
commit
2b2ea4ca37
@ -2,7 +2,8 @@
|
||||
# This is a TF Lite delegate that is powered by TensorFlow's Eager.
|
||||
#
|
||||
package(default_visibility = [
|
||||
"//visibility:private",
|
||||
"//tensorflow/compiler/mlir/lite:__subpackages__",
|
||||
"//tensorflow/lite/toco/tflite:__subpackages__",
|
||||
])
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
@ -241,3 +242,13 @@ tf_cc_test(
|
||||
"@com_google_googletest//:gtest",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "whitelisted_flex_ops_lib",
|
||||
srcs = [
|
||||
"whitelisted_flex_ops.cc",
|
||||
],
|
||||
hdrs = [
|
||||
"whitelisted_flex_ops.h",
|
||||
],
|
||||
)
|
||||
|
@ -12,12 +12,12 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
#include "tensorflow/lite/toco/tflite/whitelisted_flex_ops.h"
|
||||
#include "tensorflow/lite/delegates/flex/whitelisted_flex_ops.h"
|
||||
|
||||
#include <set>
|
||||
|
||||
namespace toco {
|
||||
namespace tflite {
|
||||
namespace flex {
|
||||
|
||||
bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
|
||||
static const std::set<std::string>* whitelisted_flex_ops =
|
||||
@ -473,5 +473,5 @@ bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name) {
|
||||
whitelisted_flex_ops->end();
|
||||
}
|
||||
|
||||
} // namespace flex
|
||||
} // namespace tflite
|
||||
} // namespace toco
|
@ -12,13 +12,13 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
#ifndef TENSORFLOW_LITE_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
|
||||
#define TENSORFLOW_LITE_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
|
||||
#ifndef TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_
|
||||
#define TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace toco {
|
||||
namespace tflite {
|
||||
namespace flex {
|
||||
|
||||
// Whether the given op has been statically whitelisted for flex export.
|
||||
//
|
||||
@ -29,7 +29,7 @@ namespace tflite {
|
||||
// TODO(b/118389105): Automate generation of the whitelisted flex ops.
|
||||
bool IsWhitelistedFlexOp(const std::string& tensorflow_op_name);
|
||||
|
||||
} // namespace flex
|
||||
} // namespace tflite
|
||||
} // namespace toco
|
||||
|
||||
#endif // TENSORFLOW_LITE_TOCO_TFLITE_WHITELISTED_FLEX_OPS_H_
|
||||
#endif // TENSORFLOW_LITE_DELEGATES_FLEX_WHITELISTED_FLEX_OPS_H_
|
@ -14,8 +14,6 @@ cc_library(
|
||||
name = "operator",
|
||||
srcs = [
|
||||
"operator.cc",
|
||||
"whitelisted_flex_ops.cc",
|
||||
"whitelisted_flex_ops.h",
|
||||
],
|
||||
hdrs = [
|
||||
"builtin_operator.h",
|
||||
@ -31,6 +29,7 @@ cc_library(
|
||||
"//tensorflow/core:framework",
|
||||
"//tensorflow/core:protos_all_cc",
|
||||
"//tensorflow/core:ptr_util",
|
||||
"//tensorflow/lite/delegates/flex:whitelisted_flex_ops_lib",
|
||||
"//tensorflow/lite/schema:schema_fbs",
|
||||
"//tensorflow/lite/toco:graph_transformations",
|
||||
"//tensorflow/lite/toco:model",
|
||||
|
@ -22,6 +22,7 @@ limitations under the License.
|
||||
|
||||
// TODO(ycling): Consider refactoring to extract the LSTM definition out of
|
||||
// graph_transformation module.
|
||||
#include "tensorflow/lite/delegates/flex/whitelisted_flex_ops.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
#include "tensorflow/lite/toco/graph_transformations/lstm_utils.h"
|
||||
#include "tensorflow/lite/toco/model.h"
|
||||
@ -29,7 +30,6 @@ limitations under the License.
|
||||
#include "tensorflow/lite/toco/tflite/custom_operator.h"
|
||||
#include "tensorflow/lite/toco/tflite/simple_operator.h"
|
||||
#include "tensorflow/lite/toco/tflite/types.h"
|
||||
#include "tensorflow/lite/toco/tflite/whitelisted_flex_ops.h"
|
||||
|
||||
namespace toco {
|
||||
|
||||
@ -2646,7 +2646,7 @@ bool ShouldExportAsFlexOp(bool enable_select_tf_ops,
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!IsWhitelistedFlexOp(tensorflow_op_name)) {
|
||||
if (!::tflite::flex::IsWhitelistedFlexOp(tensorflow_op_name)) {
|
||||
LOG(WARNING) << "Op " << tensorflow_op_name
|
||||
<< " is a valid TensorFlow op but has not been whitelisted for"
|
||||
" the TensorFlow Lite flex op set.";
|
||||
|
Loading…
x
Reference in New Issue
Block a user