Skip all ops with function attribute by default

2 files changed, 13 insertions, 2 deletions

diff --git a/tensorflow/java/BUILD b/tensorflow/java/BUILD
index 7cd0208dbf..0cc8e7c3e2 100644
--- a/tensorflow/java/BUILD
+++ b/tensorflow/java/BUILD
@@ -72,7 +72,6 @@ tf_java_op_gen_srcjar(
     name = "java_op_gen_sources",
     api_def_srcs = [
         "//tensorflow/core/api_def:base_api_def",
-        "//tensorflow/core/api_def:java_api_def",
     ],
     base_package = "org.tensorflow.op",
     gen_tool = ":java_op_gen_tool",
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index 7355b3a395..f4cefbe933 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -420,6 +420,18 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   writer.EndType();
 }
 
+bool CanGenerateOp(const OpDef& op_def, const ApiDef& api_def) {
+  if (api_def.visibility() == ApiDef::SKIP) {
+    return false;
+  }
+  for (const auto& attr : op_def.attr()) {
+    if (attr.type() == "func") {
+      return false;  // TODO(karllessard) add support for function attributes
+    }
+  }
+  return true;
+}
+
 }  // namespace
 
 Status OpGenerator::Run(const OpList& op_list, const string& base_package,
@@ -441,7 +453,7 @@ Status OpGenerator::Run(const OpList& op_list, const string& base_package,
   api_map.UpdateDocs();
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
-    if (api_def->visibility() != ApiDef::SKIP) {
+    if (CanGenerateOp(op_def, *api_def)) {
       OpSpec op(OpSpec::Create(op_def, *api_def));
       for (const EndpointSpec& endpoint : op.endpoints()) {
         GenerateOp(op, endpoint, base_package, output_dir, env_);