Update annotation names and comments, and registerCustomGradient javadoc

Signed-off-by: Ryan Nett <[email protected]>

Author: rnett

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/TensorFlow.java

@@ -39,8 +39,8 @@
 import org.tensorflow.op.CustomGradient;
 import org.tensorflow.op.RawCustomGradient;
 import org.tensorflow.op.RawOpInputs;
-import org.tensorflow.op.annotation.GeneratedOpInputsMetadata;
-import org.tensorflow.op.annotation.GeneratedOpMetadata;
+import org.tensorflow.op.annotation.OpInputsMetadata;
+import org.tensorflow.op.annotation.OpMetadata;
 import org.tensorflow.op.math.Add;
 import org.tensorflow.proto.framework.OpList;
 
@@ -187,35 +187,34 @@ public static synchronized boolean registerCustomGradient(
   }
 
   /**
-   * Register a custom gradient function for ops of {@code opClass} type. The actual op type is
-   * detected from the class's {@link GeneratedOpMetadata} annotation. As such, it only works on
-   * generated op classes.
+   * Register a custom gradient function for ops of {@code inputClass}'s op type. The actual op type
+   * is detected from the class's {@link OpInputsMetadata} annotation. As such, it only works on
+   * generated op classes or custom op classes with the correct annotations.
    *
-   * @param opClass the class of op to register the gradient for.
+   * @param inputClass the inputs class of op to register the gradient for.
    * @param gradient the gradient function to use
    * @return {@code true} if the gradient was registered, {@code false} if there was already a
    *     gradient registered for this op
-   * @throws IllegalArgumentException if {@code opClass} does not have a {@link GeneratedOpMetadata}
-   *     field.
+   * @throws IllegalArgumentException if {@code inputClass} is not annotated with {@link
+   *     OpInputsMetadata} or the op class is not annotated with {@link OpMetadata}.
    */
   public static synchronized <T extends RawOpInputs<?>> boolean registerCustomGradient(
-      Class<T> opClass, CustomGradient<T> gradient) {
-    GeneratedOpInputsMetadata metadata = opClass.getAnnotation(GeneratedOpInputsMetadata.class);
+      Class<T> inputClass, CustomGradient<T> gradient) {
+    OpInputsMetadata metadata = inputClass.getAnnotation(OpInputsMetadata.class);
 
     if (metadata == null) {
       throw new IllegalArgumentException(
           "Inputs Class "
-              + opClass
-              + " does not have a GeneratedOpInputsMetadata annotation.  Was it generated by tensorflow/java?  If it was, this is a bug.");
+              + inputClass
+              + " does not have a OpInputsMetadata annotation.  Was it generated by tensorflow/java?  If it was, this is a bug.");
     }
-    GeneratedOpMetadata outputMetadata =
-        metadata.outputsClass().getAnnotation(GeneratedOpMetadata.class);
+    OpMetadata outputMetadata = metadata.outputsClass().getAnnotation(OpMetadata.class);
 
     if (outputMetadata == null) {
       throw new IllegalArgumentException(
           "Op Class "
               + metadata.outputsClass()
-              + " does not have a GeneratedOpMetadata annotation.  Was it generated by tensorflow/java?  If it was, this is a bug.");
+              + " does not have a OpMetadata annotation.  Was it generated by tensorflow/java?  If it was, this is a bug.");
     }
 
     String opType = outputMetadata.opType();
@@ -224,7 +223,7 @@ public static synchronized <T extends RawOpInputs<?>> boolean registerCustomGrad
       return false;
     }
 
-    GradFunc g = CustomGradient.adapter(gradient, opClass);
+    GradFunc g = CustomGradient.adapter(gradient, inputClass);
     GradOpRegistry.Global().Register(opType, g);
     gradientFuncs.add(g);
     return true;

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/annotation/GeneratedOpInputsMetadata.java renamed to tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/annotation/OpInputsMetadata.java

@@ -24,13 +24,13 @@
 import org.tensorflow.op.RawOp;
 
 /**
- * An annotation that should only be used by codegeneration. Used to provide some metadata about the
- * op.
- *
- * <p><b>DO NOT USE MANUALLY</b>
+ * An annotation to provide metadata about an op inputs accessor class. Should only be used by users
+ * on custom ops, will be generated for non-custom ops.
  */
 @Target(ElementType.TYPE)
 @Retention(RetentionPolicy.RUNTIME)
-public @interface GeneratedOpInputsMetadata {
+public @interface OpInputsMetadata {
+
+  /** The main op class. */
   Class<? extends RawOp> outputsClass();
 }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/annotation/GeneratedOpMetadata.java renamed to tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/annotation/OpMetadata.java

@@ -24,16 +24,17 @@
 import org.tensorflow.op.RawOpInputs;
 
 /**
- * An annotation that should only be used by codegeneration. Used to provide some metadata about the
- * op.
- *
- * <p><b>DO NOT USE MANUALLY</b>
+ * An annotation to provide metadata about an op. Should only be used by users on custom ops, will
+ * be generated for non-custom ops.
  */
 @Target(ElementType.TYPE)
 @Retention(RetentionPolicy.RUNTIME)
-public @interface GeneratedOpMetadata {
+public @interface OpMetadata {
+
+  /** The type of the op in the TF runtime. */
   String opType();
 
+  /** The typesafe inputs class (which should be annotated with {@link OpInputsMetadata}). */
   @SuppressWarnings("rawtypes")
   Class<? extends RawOpInputs> inputsClass();
 }

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/Names.java

@@ -29,10 +29,10 @@ public class Names {
 
   public static final ClassName Operator = ClassName.get(OpPackage + ".annotation", "Operator");
   public static final ClassName Endpoint = ClassName.get(OpPackage + ".annotation", "Endpoint");
-  public static final ClassName GeneratedOpMetadata =
-      ClassName.get(OpPackage + ".annotation", "GeneratedOpMetadata");
-  public static final ClassName GeneratedOpInputsMetadata =
-      ClassName.get(OpPackage + ".annotation", "GeneratedOpInputsMetadata");
+  public static final ClassName OpMetadata =
+      ClassName.get(OpPackage + ".annotation", "OpMetadata");
+  public static final ClassName OpInputsMetadata =
+      ClassName.get(OpPackage + ".annotation", "OpInputsMetadata");
 
   public static final ClassName TType = ClassName.get(TypesPackage + ".family", "TType");
   public static final ClassName TString = ClassName.get(TypesPackage, "TString");

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/op/generator/ClassGenerator.java

@@ -1057,7 +1057,7 @@ private Set<TypeVariableName> buildInputsClass() {
   /** Adds the GeneratedOpMetadata annotation to the op class. */
   private void addInputsMetadataAnnotation() {
     builder.addAnnotation(
-        AnnotationSpec.builder(Names.GeneratedOpMetadata)
+        AnnotationSpec.builder(Names.OpMetadata)
             .addMember("opType", "$L", className + ".OP_NAME")
             .addMember("inputsClass", "$T.class", inputsClassName())
             .build());
@@ -1066,7 +1066,7 @@ private void addInputsMetadataAnnotation() {
   /** Adds the GeneratedOpInputsMetadata annotation to the op input class. */
   private void addInputsMetadataAnnotation(TypeSpec.Builder inputsBuilder) {
     inputsBuilder.addAnnotation(
-        AnnotationSpec.builder(Names.GeneratedOpInputsMetadata)
+        AnnotationSpec.builder(Names.OpInputsMetadata)
             .addMember("outputsClass", "$T.class", className())
             .build());
   }