Use annotation instead of field reflection to store op types

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

Author: rnett

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

@@ -23,8 +23,6 @@
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_Version;
 
 import com.google.protobuf.InvalidProtocolBufferException;
-import java.lang.reflect.Field;
-import java.lang.reflect.Modifier;
 import java.util.Collections;
 import java.util.IdentityHashMap;
 import java.util.Set;
@@ -43,6 +41,7 @@
 import org.tensorflow.op.RawGradientAdapter;
 import org.tensorflow.op.RawOp;
 import org.tensorflow.op.TypedGradientAdapter;
+import org.tensorflow.op.annotation.GeneratedOpMetadata;
 import org.tensorflow.op.math.Add;
 import org.tensorflow.proto.framework.OpList;
 
@@ -190,41 +189,36 @@ 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 {@code OP_NAME} field.
+   * detected from the class's {@link GeneratedOpMetadata} annotation. As such, it only works on
+   * generated op classes.
    *
    * @param opClass the 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 static {@code OP_NAME}
+   * @throws IllegalArgumentException if {@code opClass} does not have a {@link GeneratedOpMetadata}
    *     field.
    */
   public static synchronized <T extends RawOp> boolean registerCustomGradient(
       Class<T> opClass, CustomGradient<T> gradient) {
-    try {
-      Field nameField = opClass.getDeclaredField("OP_NAME");
-
-      if (!Modifier.isStatic(nameField.getModifiers())) {
-        throw new IllegalArgumentException(
-            "Class " + opClass + " has an OP_NAME field, but it is not static.");
-      }
-
-      String opType = (String) nameField.get(null);
+    GeneratedOpMetadata metadata = opClass.getAnnotation(GeneratedOpMetadata.class);
 
-      if (hasGradient(opType)) {
-        return false;
-      }
-
-      GradFunc g = new TypedGradientAdapter<>(gradient, opClass);
-      GradOpRegistry.Global().Register(opType, g);
-      gradientFuncs.add(g);
-    } catch (IllegalAccessException | NoSuchFieldException e) {
+    if (metadata == null) {
       throw new IllegalArgumentException(
-          "Could not get OP_NAME field for class "
+          "Class "
               + opClass
-              + ", ensure it is a generated op class",
-          e);
+              + " does not have a GeneratedOpMetadata annotation.  Was it generated by tensorflow/java?  If it was, this is a bug.");
+    }
+
+    String opType = metadata.opType();
+
+    if (hasGradient(opType)) {
+      return false;
     }
+
+    GradFunc g = new TypedGradientAdapter<>(gradient, opClass);
+    GradOpRegistry.Global().Register(opType, g);
+    gradientFuncs.add(g);
     return true;
   }
 }

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

@@ -0,0 +1,38 @@
+/*
+ Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+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.
+=======================================================================
+
+*/
+package org.tensorflow.op.annotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.function.Function;
+import org.tensorflow.Operation;
+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>
+ */
+@Target(ElementType.TYPE)
+@Retention(RetentionPolicy.RUNTIME)
+public @interface GeneratedOpMetadata {
+  String opType();
+}

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

@@ -29,6 +29,8 @@ 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 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

@@ -220,6 +220,7 @@ void buildClass() {
     if (!isStateSelector) {
       builder.addModifiers(Modifier.FINAL);
       builder.superclass(Names.RawOp);
+      addMetadataAnnotation();
     }
 
     if (isStateSubclass) {
@@ -1031,4 +1032,12 @@ private void buildInputsClass() {
     inputsBuilder.addTypeVariables(typeVars);
     this.builder.addType(inputsBuilder.build());
   }
+
+  /** Adds the GeneratedOpMetadata annotation */
+  private void addMetadataAnnotation() {
+    builder.addAnnotation(
+        AnnotationSpec.builder(Names.GeneratedOpMetadata)
+            .addMember("opType", "$L", className + ".OP_NAME")
+            .build());
+  }
 }

tensorflow-framework/src/main/java/org/tensorflow/framework/op/FrameworkOps.java

@@ -17,8 +17,8 @@
 import org.tensorflow.DeviceSpec;
 import org.tensorflow.EagerSession;
 import org.tensorflow.ExecutionEnvironment;
-import org.tensorflow.op.OpScope;
 import org.tensorflow.op.Op;
+import org.tensorflow.op.OpScope;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.Scope;
 

tensorflow-framework/src/main/java/org/tensorflow/framework/optimizers/Optimizer.java

@@ -25,8 +25,8 @@
 import org.tensorflow.Operand;
 import org.tensorflow.Operation;
 import org.tensorflow.Output;
-import org.tensorflow.op.OpScope;
 import org.tensorflow.op.Op;
+import org.tensorflow.op.OpScope;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.Scope;
 import org.tensorflow.op.core.NoOp;