Simplify and improve generics handling in generator

6 files changed, 132 insertions, 138 deletions

diff --git a/tensorflow/java/build_defs.bzl b/tensorflow/java/build_defs.bzl
index ab7f60d03d..e1916ca4d9 100644
--- a/tensorflow/java/build_defs.bzl
+++ b/tensorflow/java/build_defs.bzl
@@ -15,6 +15,7 @@ JAVA_VERSION_OPTS = [
 XLINT_OPTS = [
     "-Werror",
     "-Xlint:all",
+    "-Xlint:-processing",
     "-Xlint:-serial",
     "-Xlint:-try",
     "-Xlint:-classfile", # see b/32750402, go/javac-warnings#classfile
diff --git a/tensorflow/java/src/gen/cc/op_gen_main.cc b/tensorflow/java/src/gen/cc/op_gen_main.cc
index 458141b877..a508c96516 100644
--- a/tensorflow/java/src/gen/cc/op_gen_main.cc
+++ b/tensorflow/java/src/gen/cc/op_gen_main.cc
@@ -67,10 +67,10 @@ int main(int argc, char* argv[]) {
   QCHECK(parsed_flags_ok && !output_dir.empty()) << usage;
   std::vector<tensorflow::string> api_dirs = tensorflow::str_util::Split(
       api_dirs_str, ",", tensorflow::str_util::SkipEmpty());
-  tensorflow::java::OpGenerator generator(base_package, output_dir, api_dirs);
+  tensorflow::java::OpGenerator generator(api_dirs);
   tensorflow::OpList ops;
   tensorflow::OpRegistry::Global()->Export(false, &ops);
-  TF_CHECK_OK(generator.Run(ops));
+  TF_CHECK_OK(generator.Run(ops, base_package, output_dir));
 
   return 0;
 }
diff --git a/tensorflow/java/src/gen/cc/op_generator.cc b/tensorflow/java/src/gen/cc/op_generator.cc
index 00f84bc9cd..2327a4daf1 100644
--- a/tensorflow/java/src/gen/cc/op_generator.cc
+++ b/tensorflow/java/src/gen/cc/op_generator.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <list>
 #include <memory>
 #include <set>
+#include <ctime>
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/str_util.h"
@@ -38,23 +39,18 @@ namespace {
 const char* kLicenseSnippet =
     "tensorflow/java/src/gen/resources/license.java.snippet";
 
-const std::map<string, Type> kPrimitiveAttrTypes = {
-  { "Boolean", Type::Boolean() },
-  { "Byte", Type::Byte() },
-  { "Character", Type::Byte() },
-  { "Float", Type::Float() },
-  { "Integer", Type::Long() },
-  { "Long", Type::Long() },
-  { "Short", Type::Long() },
-  { "Double", Type::Float() },
-};
-
 enum RenderMode {
   DEFAULT,
   SINGLE_OUTPUT,
   SINGLE_LIST_OUTPUT
 };
 
+inline void AddArgument(const Variable& var, const string& description,
+    Method* method_out, Javadoc* javadoc_out) {
+  method_out->add_argument(var);
+  javadoc_out->add_param_tag(var.name(), description);
+}
+
 void CollectOpDependencies(const OpSpec& op, RenderMode mode,
     std::list<Type>* out) {
   out->push_back(Type::Class("Operation", "org.tensorflow"));
@@ -81,9 +77,7 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
   }
   for (const AttributeSpec& attribute : op.attributes()) {
     out->push_back(attribute.var().type());
-    if (attribute.var().type().name() == "Class") {
-      out->push_back(Type::Enum("DataType", "org.tensorflow"));
-    }
+    out->push_back(attribute.jni_type());
   }
   for (const AttributeSpec& optional_attribute : op.optional_attributes()) {
     out->push_back(optional_attribute.var().type());
@@ -92,45 +86,38 @@ void CollectOpDependencies(const OpSpec& op, RenderMode mode,
 
 void WriteSetAttrDirective(const AttributeSpec& attr, bool optional,
     SourceWriter* writer) {
-  string var = optional ? "opts." + attr.var().name() : attr.var().name();
+  string var_name = optional ? "opts." + attr.var().name() : attr.var().name();
   if (attr.iterable()) {
-    const Type& type = attr.type();
-    std::map<string, Type>::const_iterator it =
-      kPrimitiveAttrTypes.find(type.name());
-    if (it != kPrimitiveAttrTypes.end()) {
-      string array = attr.var().name() + "Array";
-      writer->AppendType(it->second)
-          .Append("[] " + array + " = new ")
-          .AppendType(it->second)
-          .Append("[" + var + ".size()];")
-          .EndLine();
-      writer->BeginBlock("for (int i = 0; i < " + array + ".length; ++i)")
-          .Append(array + "[i] = " + var + ".get(i);")
-          .EndLine()
-          .EndBlock()
-          .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + array)
-          .Append(");")
-          .EndLine();
+    string array_name = attr.var().name() + "Array";
+    writer->AppendType(attr.jni_type())
+        .Append("[] " + array_name + " = new ")
+        .AppendType(attr.jni_type())
+        .Append("[" + var_name + ".size()];")
+        .EndLine()
+        .BeginBlock("for (int i = 0; i < " + array_name + ".length; ++i)")
+        .Append(array_name + "[i] = ");
+    if (attr.type().kind() == Type::GENERIC) {
+      writer->Append("DataType.fromClass(" + var_name + ".get(i));");
     } else {
-      writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", " + var)
-          .Append(".toArray(new ")
-          .AppendType(type)
-          .Append("[" + var + ".size()]));")
-          .EndLine();
+      writer->Append(var_name + ".get(i);");
     }
+    writer->EndLine()
+        .EndBlock()
+        .Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ")
+        .Append(array_name + ");")
+        .EndLine();
   } else {
-    Type type = attr.var().type();
     writer->Append("opBuilder.setAttr(\"" + attr.op_def_name() + "\", ");
-    if (type.name() == "Class") {
-      writer->Append("DataType.fromClass(" + attr.var().name() + "));");
+    if (attr.var().type().name() == "Class") {
+      writer->Append("DataType.fromClass(" + var_name + "));");
     } else {
-      writer->Append(var + ");");
+      writer->Append(var_name + ");");
     }
     writer->EndLine();
   }
 }
 
-void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
+void RenderFactoryMethods(const OpSpec& op, const Type& op_class,
     SourceWriter* writer) {
   Method factory = Method::Create("create", op_class);
   Javadoc factory_doc = Javadoc::Create(
@@ -138,27 +125,24 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
       + " operation to the graph.");
   Variable scope =
       Variable::Create("scope", Type::Class("Scope", "org.tensorflow.op"));
-  factory.add_argument(scope);
-  factory_doc.add_param_tag(scope.name(), "Current graph scope");
+  AddArgument(scope, "current graph scope", &factory, &factory_doc);
   for (const ArgumentSpec& input : op.inputs()) {
-    factory.add_argument(input.var());
-    factory_doc.add_param_tag(input.var().name(), input.description());
+    AddArgument(input.var(), input.description(), &factory, &factory_doc);
   }
-  for (const AttributeSpec& attribute : op.attributes()) {
-    factory.add_argument(attribute.var());
-    factory_doc.add_param_tag(attribute.var().name(), attribute.description());
+  for (const AttributeSpec& attr : op.attributes()) {
+    AddArgument(attr.var(), attr.description(), &factory, &factory_doc);
   }
   if (!op.optional_attributes().empty()) {
-    factory.add_argument(Variable::Varargs("options", Type::Class("Options")));
-    factory_doc.add_param_tag("options", "carries optional attributes values");
+    AddArgument(Variable::Varargs("options", Type::Class("Options")),
+        "carries optional attributes values", &factory, &factory_doc);
   }
   factory_doc.add_tag("return", "a new instance of " + op_class.name());
+
   writer->BeginMethod(factory, PUBLIC|STATIC, &factory_doc);
   writer->Append("OperationBuilder opBuilder = scope.graph().opBuilder(\""
       + op.graph_op_name() + "\", scope.makeOpName(\""
       + op_class.name() + "\"));");
   writer->EndLine();
-
   for (const ArgumentSpec& input : op.inputs()) {
     if (input.iterable()) {
       writer->Append("opBuilder.addInputList(Operands.asOutputs("
@@ -192,10 +176,9 @@ void RenderFactoryMethod(const OpSpec& op, const Type& op_class,
 
 void RenderConstructor(const OpSpec& op, const Type& op_class,
     SourceWriter* writer) {
-  Method constructor = Method::ConstructorFor(op_class)
-    .add_argument(
-        Variable::Create("operation",
-            Type::Class("Operation", "org.tensorflow")));
+  Variable operation =
+      Variable::Create("operation", Type::Class("Operation", "org.tensorflow"));
+  Method constructor = Method::ConstructorFor(op_class).add_argument(operation);
   for (const ArgumentSpec& output : op.outputs()) {
     if (output.iterable() && !output.type().unknown()) {
       constructor.add_annotation(
@@ -237,15 +220,14 @@ void RenderConstructor(const OpSpec& op, const Type& op_class,
 }
 
 void RenderGettersAndSetters(const OpSpec& op, SourceWriter* writer) {
-  for (const AttributeSpec& attribute : op.optional_attributes()) {
+  for (const AttributeSpec& attr : op.optional_attributes()) {
     Method setter =
-        Method::Create(attribute.var().name(), Type::Class("Options"))
-            .add_argument(attribute.var());
-    Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(attribute.var().name(), attribute.description());
+        Method::Create(attr.var().name(), Type::Class("Options"));
+    Javadoc setter_doc = Javadoc::Create();
+    AddArgument(attr.var(), attr.description(), &setter, &setter_doc);
     writer->BeginMethod(setter, PUBLIC|STATIC, &setter_doc)
-        .Append("return new Options()." + attribute.var().name() + "("
-            + attribute.var().name() + ");")
+        .Append("return new Options()." + attr.var().name() + "("
+            + attr.var().name() + ");")
         .EndLine()
         .EndMethod();
   }
@@ -311,14 +293,12 @@ void RenderOptionsClass(const OpSpec& op, const Type& op_class,
   Javadoc options_doc = Javadoc::Create(
       "Optional attributes for {@link " + op_class.full_name() + "}");
   writer->BeginInnerType(options_class, PUBLIC | STATIC, &options_doc);
-  for (const AttributeSpec& attribute : op.optional_attributes()) {
-    Method setter = Method::Create(attribute.var().name(), options_class)
-        .add_argument(attribute.var());
-    Javadoc setter_doc = Javadoc::Create()
-        .add_param_tag(attribute.var().name(), attribute.description());
+  for (const AttributeSpec& attr : op.optional_attributes()) {
+    Method setter = Method::Create(attr.var().name(), options_class);
+    Javadoc setter_doc = Javadoc::Create();
+    AddArgument(attr.var(), attr.description(), &setter, &setter_doc);
     writer->BeginMethod(setter, PUBLIC, &setter_doc)
-        .Append("this." + attribute.var().name() + " = "
-            + attribute.var().name() + ";")
+        .Append("this." + attr.var().name() + " = " + attr.var().name() + ";")
         .EndLine()
         .Append("return this;")
         .EndLine()
@@ -339,12 +319,13 @@ inline Type ClassOf(const EndpointSpec& endpoint, const string& base_package) {
 }
 
 void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
-    const string& base_package, const string& output_dir, Env* env) {
+    const string& base_package, const string& output_dir, Env* env,
+    const std::tm* timestamp) {
   Type op_class(ClassOf(endpoint, base_package)
       .add_supertype(Type::Class("PrimitiveOp", "org.tensorflow.op")));
   Javadoc op_javadoc(endpoint.javadoc());
 
-  // implement Operand (or Iterable<Operand>) if the op has only one output
+  // op interfaces
   RenderMode mode = DEFAULT;
   if (op.outputs().size() == 1) {
     const ArgumentSpec& output = op.outputs().front();
@@ -360,18 +341,22 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
       op_class.add_supertype(operand_inf);
     }
   }
-  // declare all outputs generics at the op class level
+  // op generic parameters
   std::set<string> generics;
   for (const ArgumentSpec& output : op.outputs()) {
     if (output.type().kind() == Type::GENERIC && !output.type().unknown()
         && generics.find(output.type().name()) == generics.end()) {
       op_class.add_parameter(output.type());
       op_javadoc.add_param_tag("<" + output.type().name() + ">",
-          "data type of output {@code " + output.var().name() + "}");
+          "data type for {@code " + output.var().name() + "()} output");
       generics.insert(output.type().name());
     }
   }
-  // handle endpoint deprecation
+  // op annotations
+  char date[20];
+  strftime(date, sizeof date, "%FT%TZ", timestamp);
+  op_class.add_annotation(Annotation::Create("Generated", "javax.annotation")
+      .attributes(string("value = \"op_generator\", date = \"") + date + "\""));
   if (endpoint.deprecated()) {
     op_class.add_annotation(Annotation::Create("Deprecated"));
     string explanation;
@@ -384,8 +369,8 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
     }
     op_javadoc.add_tag("deprecated", explanation);
   }
-  // expose the op in the Ops Graph API only if it is visible
   if (!op.hidden()) {
+    // expose the op in the Ops Graph API only if it is visible
     op_class.add_annotation(
         Annotation::Create("Operator", "org.tensorflow.op.annotation")
           .attributes("group = \"" + endpoint.package() + "\""));
@@ -406,14 +391,11 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
   CollectOpDependencies(op, mode, &dependencies);
   writer.WriteFromFile(kLicenseSnippet)
       .EndLine()
-      .Append("// This file is machine generated, DO NOT EDIT!")
-      .EndLine()
-      .EndLine()
       .BeginType(op_class, PUBLIC|FINAL, &dependencies, &op_javadoc);
   if (!op.optional_attributes().empty()) {
     RenderOptionsClass(op, op_class, &writer);
   }
-  RenderFactoryMethod(op, op_class, &writer);
+  RenderFactoryMethods(op, op_class, &writer);
   RenderGettersAndSetters(op, &writer);
   if (mode != DEFAULT) {
     RenderInterfaceImpl(op, mode, &writer);
@@ -428,13 +410,8 @@ void GenerateOp(const OpSpec& op, const EndpointSpec& endpoint,
 
 }  // namespace
 
-OpGenerator::OpGenerator(const string& base_package, const string& output_dir,
-    const std::vector<string>& api_dirs, Env* env)
-  : base_package_(base_package), output_dir_(output_dir), api_dirs_(api_dirs),
-    env_(env) {
-}
-
-Status OpGenerator::Run(const OpList& op_list) {
+Status OpGenerator::Run(const OpList& op_list, const string& base_package,
+    const string& output_dir) {
   ApiDefMap api_map(op_list);
   if (!api_dirs_.empty()) {
     // Only load api files that correspond to the requested "op_list"
@@ -449,12 +426,14 @@ Status OpGenerator::Run(const OpList& op_list) {
     }
   }
   api_map.UpdateDocs();
+  time_t now;
+  time(&now);
   for (const auto& op_def : op_list.op()) {
     const ApiDef* api_def = api_map.GetApiDef(op_def.name());
     if (api_def->visibility() != ApiDef::SKIP) {
       OpSpec op(OpSpec::Create(op_def, *api_def));
       for (const EndpointSpec& endpoint : op.endpoints()) {
-        GenerateOp(op, endpoint, base_package_, output_dir_, env_);
+        GenerateOp(op, endpoint, base_package, output_dir, env_, gmtime(&now));
       }
     }
   }
diff --git a/tensorflow/java/src/gen/cc/op_generator.h b/tensorflow/java/src/gen/cc/op_generator.h
index 06b08e852a..b789e11fa9 100644
--- a/tensorflow/java/src/gen/cc/op_generator.h
+++ b/tensorflow/java/src/gen/cc/op_generator.h
@@ -36,18 +36,17 @@ namespace java {
 // ops definitions.
 class OpGenerator {
  public:
-  OpGenerator(const string& base_package, const string& output_dir,
-      const std::vector<string>& api_dirs, Env* env = Env::Default());
+  explicit OpGenerator(const std::vector<string>& api_dirs,
+      Env* env = Env::Default()) : api_dirs_(api_dirs), env_(env) {}
 
   // Generates wrappers for the given list of 'ops'.
   //
-  // Output files are generated in <output_dir>/<base_package>/<lib_package>,
-  // where 'lib_package' is derived from ops endpoints.
-  Status Run(const OpList& op_list);
+  // Output files are generated in <output_dir>/<base_package>/<op_package>,
+  // where 'op_package' is derived from ops endpoints.
+  Status Run(const OpList& op_list, const string& base_package,
+      const string& output_dir);
 
  private:
-  const string base_package_;
-  const string output_dir_;
   const std::vector<string> api_dirs_;
   Env* env_;
 };
diff --git a/tensorflow/java/src/gen/cc/op_specs.cc b/tensorflow/java/src/gen/cc/op_specs.cc
index a0e7a180f2..dcc6388614 100644
--- a/tensorflow/java/src/gen/cc/op_specs.cc
+++ b/tensorflow/java/src/gen/cc/op_specs.cc
@@ -46,14 +46,30 @@ class TypeResolver {
   explicit TypeResolver(const OpDef& op_def) : op_def_(op_def) {}
 
   Type TypeOf(const OpDef_ArgDef& arg_def, bool *iterable_out);
-  Type TypeOf(const OpDef_AttrDef& attr_def, bool *iterable_out);
+  std::pair<Type, Type> TypeOf(const OpDef_AttrDef& attr_def,
+      bool *iterable_out);
   bool IsAttributeVisited(const string& attr_name) {
     return visited_attrs_.find(attr_name) != visited_attrs_.cend();
   }
+
  private:
   const OpDef op_def_;
   std::map<std::string, Type> visited_attrs_;
-  char next_generic_ = 'T';
+  char next_generic_letter_ = 'T';
+
+  std::pair<Type, Type> MakeTypePair(const Type& type, const Type& jni_type) {
+    return std::make_pair(type, jni_type);
+  }
+  std::pair<Type, Type> MakeTypePair(const Type& type) {
+    return std::make_pair(type, type);
+  }
+  Type NextGeneric() {
+    char generic_letter = next_generic_letter_++;
+    if (next_generic_letter_ > 'Z') {
+      next_generic_letter_ = 'A';
+    }
+    return Type::Generic(string(1, generic_letter));
+  }
 };
 
 Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
@@ -107,7 +123,7 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
     } else {
       for (const auto& attr_def : op_def_.attr()) {
         if (attr_def.name() == arg_def.type_attr()) {
-          type = TypeOf(attr_def, iterable_out);
+          type = TypeOf(attr_def, iterable_out).first;
           break;
         }
       }
@@ -125,51 +141,47 @@ Type TypeResolver::TypeOf(const OpDef_ArgDef& arg_def,
   return type;
 }
 
-Type TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
+std::pair<Type, Type> TypeResolver::TypeOf(const OpDef_AttrDef& attr_def,
     bool* iterable_out) {
+  std::pair<Type, Type> types = MakeTypePair(Type::Wildcard());
   *iterable_out = false;
   StringPiece attr_type = attr_def.type();
   if (str_util::ConsumePrefix(&attr_type, "list(")) {
     attr_type.remove_suffix(1);  // remove closing brace
     *iterable_out = true;
   }
-  Type type = *iterable_out ? Type::Wildcard() : Type::Class("Object");
-  if (attr_type == "type") {
-    if (*iterable_out) {
-      type = Type::Enum("DataType", "org.tensorflow");
-    } else {
-      type = Type::Generic(string(1, next_generic_));
-      next_generic_ = (next_generic_ == 'Z') ? 'A' : next_generic_ + 1;
-      if (IsRealNumbers(attr_def.allowed_values())) {
-        // enforce real numbers datasets by extending java.lang.Number
-        type.add_supertype(Type::Class("Number"));
-      }
-    }
-  } else if (attr_type == "string") {
-    type = Type::Class("String");
+  if (attr_type == "string") {
+    types = MakeTypePair(Type::Class("String"));
 
   } else if (attr_type == "int") {
-    type = Type::Class("Integer");
+    types = MakeTypePair(Type::Class("Long"), Type::Long());
 
   } else if (attr_type == "float") {
-    type = Type::Class("Float");
+    types = MakeTypePair(Type::Class("Float"), Type::Float());
 
   } else if (attr_type == "bool") {
-    type = Type::Class("Boolean");
+    types = MakeTypePair(Type::Class("Boolean"), Type::Boolean());
 
   } else if (attr_type == "shape") {
-    type = Type::Class("Shape", "org.tensorflow");
+    types = MakeTypePair(Type::Class("Shape", "org.tensorflow"));
 
   } else if (attr_type == "tensor") {
-    type = Type::Class("Tensor", "org.tensorflow")
-        .add_parameter(Type::Wildcard());
+    types = MakeTypePair(Type::Class("Tensor", "org.tensorflow")
+        .add_parameter(Type::Wildcard()));
+
+  } else if (attr_type == "type") {
+    Type type = *iterable_out ? Type::Wildcard() : NextGeneric();
+    if (IsRealNumbers(attr_def.allowed_values())) {
+      type.add_supertype(Type::Class("Number"));
+    }
+    types = MakeTypePair(type, Type::Enum("DataType", "org.tensorflow"));
 
   } else {
     LOG(FATAL) << "Cannot resolve data type for attribute \"" << attr_type
         << "\" in operation \"" << op_def_.name() << "\"";
   }
-  visited_attrs_.insert(std::make_pair(attr_def.name(), type));
-  return type;
+  visited_attrs_.insert(std::make_pair(attr_def.name(), types.first));
+  return types;
 }
 
 string SnakeToCamelCase(const string& str, bool upper = false) {
@@ -307,19 +319,19 @@ ArgumentSpec CreateInput(const OpDef_ArgDef& input_def,
 AttributeSpec CreateAttribute(const OpDef_AttrDef& attr_def,
     const ApiDef::Attr& attr_api_def, TypeResolver* type_resolver) {
   bool iterable = false;
-  Type type = type_resolver->TypeOf(attr_def, &iterable);
-  // type attributes must be passed explicitly in methods as a Class<> parameter
-  bool is_explicit = type.kind() == Type::GENERIC && !iterable;
-  Type var_type = is_explicit ? Type::Class("Class").add_parameter(type) : type;
+  std::pair<Type, Type> types = type_resolver->TypeOf(attr_def, &iterable);
+  Type var_type = types.first.kind() == Type::GENERIC ?
+      Type::Class("Class").add_parameter(types.first) : types.first;
   if (iterable) {
-    var_type = Type::ListOf(type);
+    var_type = Type::ListOf(var_type);
   }
   return AttributeSpec(attr_api_def.name(),
       Variable::Create(SnakeToCamelCase(attr_api_def.rename_to()), var_type),
-      type,
+      types.first,
+      types.second,
       ParseDocumentation(attr_api_def.description()),
       iterable,
-      attr_api_def.has_default_value() && !is_explicit);
+      attr_api_def.has_default_value());
 }
 
 ArgumentSpec CreateOutput(const OpDef_ArgDef& output_def,
@@ -340,7 +352,6 @@ ArgumentSpec CreateOutput(const OpDef_ArgDef& output_def,
 
 EndpointSpec CreateEndpoint(const OpDef& op_def, const ApiDef& api_def,
     const ApiDef_Endpoint& endpoint_def) {
-
   std::vector<string> name_tokens = str_util::Split(endpoint_def.name(), ".");
   string package;
   string name;
@@ -381,7 +392,7 @@ OpSpec OpSpec::Create(const OpDef& op_def, const ApiDef& api_def) {
       AttributeSpec attr = CreateAttribute(op_def.attr(i), api_def.attr(i),
           &type_resolver);
       // attributes with a default value are optional
-      if (attr.optional()) {
+      if (attr.has_default_value() && attr.type().kind() != Type::GENERIC) {
         op.optional_attributes_.push_back(attr);
       } else {
         op.attributes_.push_back(attr);
diff --git a/tensorflow/java/src/gen/cc/op_specs.h b/tensorflow/java/src/gen/cc/op_specs.h
index 55c2c3f307..7d64391446 100644
--- a/tensorflow/java/src/gen/cc/op_specs.h
+++ b/tensorflow/java/src/gen/cc/op_specs.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "tensorflow/core/framework/op_def.pb.h"
 #include "tensorflow/core/framework/api_def.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/java/src/gen/cc/java_defs.h"
 
 namespace tensorflow {
@@ -87,20 +88,23 @@ class AttributeSpec : public ArgumentSpec {
   // op_def_name: attribute name, as known by TensorFlow core
   // var: a variable to represent this attribute in Java
   // type: the type of this attribute
+  // jni_type: the type of this attribute in JNI layer (see OperationBuilder)
   // description: a description of this attribute, in javadoc
   // iterable: true if this attribute is a list
-  // optional: true if this attribute does not require to be set explicitly
+  // has_default_value: true if this attribute has a default value if not set
   AttributeSpec(const string& op_def_name, const Variable& var,
-      const Type& type, const string& description, bool iterable,
-      bool optional)
+      const Type& type, const Type& jni_type, const string& description,
+      bool iterable, bool has_default_value)
     : ArgumentSpec(op_def_name, var, type, description, iterable),
-      optional_(optional) {}
+      jni_type_(jni_type), has_default_value_(has_default_value) {}
   virtual ~AttributeSpec() = default;
 
-  bool optional() const { return optional_; }
+  const Type& jni_type() const { return jni_type_; }
+  bool has_default_value() const { return has_default_value_; }
 
  private:
-  const bool optional_;
+  const Type jni_type_;
+  const bool has_default_value_;
 };
 
 class OpSpec {