Nicely handle pre-existing gradients

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

Author: rnett

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

@@ -23,14 +23,18 @@
 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;
 import java.util.stream.Collectors;
+import org.bytedeco.javacpp.PointerPointer;
 import org.bytedeco.javacpp.PointerScope;
 import org.tensorflow.exceptions.TensorFlowException;
 import org.tensorflow.internal.c_api.GradFunc;
 import org.tensorflow.internal.c_api.GradOpRegistry;
+import org.tensorflow.internal.c_api.NativeStatus;
 import org.tensorflow.internal.c_api.TF_Buffer;
 import org.tensorflow.internal.c_api.TF_Library;
 import org.tensorflow.internal.c_api.TF_Status;
@@ -150,7 +154,16 @@ private TensorFlow() {}
   }
 
   // to keep them from getting GC'd
-  private static Set<GradFunc> gradientFuncs = Collections.newSetFromMap(new IdentityHashMap<>());
+  private static final Set<GradFunc> gradientFuncs =
+      Collections.newSetFromMap(new IdentityHashMap<>());
+
+  private static synchronized boolean hasGradient(String opType) {
+    try (PointerScope scope = new PointerScope()) {
+      NativeStatus status =
+          GradOpRegistry.Global().Lookup(opType, new GradFunc(new PointerPointer<>(1)));
+      return status.ok();
+    }
+  }
 
   /**
    * Register a custom gradient function for ops of {@code opType} type.
@@ -161,12 +174,18 @@ private TensorFlow() {}
    * @param opType the type of op to register the gradient for. Should usually be an {@code OP_NAME}
    *     field, i.e. {@link Add#OP_NAME}.
    * @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
    */
-  public static synchronized void registerCustomGradient(
+  public static synchronized boolean registerCustomGradient(
       String opType, RawCustomGradient gradient) {
+    if (hasGradient(opType)) {
+      return false;
+    }
     GradFunc g = new RawGradientAdapter(gradient);
     GradOpRegistry.Global().Register(opType, g);
     gradientFuncs.add(g);
+    return true;
   }
 
   /**
@@ -175,13 +194,29 @@ public static synchronized void registerCustomGradient(
    *
    * @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}
+   *     field.
    */
-  public static synchronized <T extends RawOp> void registerCustomGradient(
+  public static synchronized <T extends RawOp> boolean registerCustomGradient(
       Class<T> opClass, CustomGradient<T> gradient) {
     try {
-      String opName = (String) opClass.getDeclaredField("OP_NAME").get(null);
+      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);
+
+      if (hasGradient(opType)) {
+        return false;
+      }
+
       GradFunc g = new TypedGradientAdapter<>(gradient, opClass);
-      GradOpRegistry.Global().Register(opName, g);
+      GradOpRegistry.Global().Register(opType, g);
       gradientFuncs.add(g);
     } catch (IllegalAccessException | NoSuchFieldException e) {
       throw new IllegalArgumentException(
@@ -190,5 +225,6 @@ public static synchronized <T extends RawOp> void registerCustomGradient(
               + ", ensure it is a generated op class",
           e);
     }
+    return true;
   }
 }

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/CustomGradientTest.java

@@ -17,54 +17,58 @@
 package org.tensorflow;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 import org.tensorflow.ndarray.index.Indices;
 import org.tensorflow.op.Ops;
-import org.tensorflow.op.core.Concat;
+import org.tensorflow.op.dtypes.Cast;
+import org.tensorflow.op.nn.NthElement;
 import org.tensorflow.proto.framework.DataType;
 import org.tensorflow.types.TFloat32;
 
 public class CustomGradientTest {
 
   @Test
-  public void testCustomConcat() {
+  public void testAlreadyExisting() {
+    assertFalse(
+        TensorFlow.registerCustomGradient(
+            Cast.class,
+            (tf, op, gradInputs) -> {
+              Operand<?> out = gradInputs.get(0);
+              Operand<?> a = tf.stridedSlice(out, Indices.slice(0, 1));
+              Operand<?> b = tf.stridedSlice(out, Indices.slice(1, 2));
+              return Arrays.asList(a, b, tf.constant(0f));
+            }));
+  }
+
+  @Test
+  public void testCustomGradient() {
     try (Graph g = new Graph();
         Session s = new Session(g)) {
-
-      TensorFlow.registerCustomGradient(
-          Concat.class,
-          (tf, op, gradInputs) -> {
-            Operand<?> out = gradInputs.get(0);
-            Operand<?> a = tf.stridedSlice(out, Indices.slice(0, 1));
-            Operand<?> b = tf.stridedSlice(out, Indices.slice(1, 2));
-            return Arrays.asList(a, b, tf.constant(0f));
-          });
+      assertTrue(
+          TensorFlow.registerCustomGradient(
+              NthElement.class,
+              (tf, op, gradInputs) -> Arrays.asList(tf.constant(0f), tf.constant(0f))));
 
       Ops tf = Ops.create(g);
+      Output<TFloat32> x = tf.placeholder(TFloat32.class).output();
+      Output<TFloat32> y = tf.nn.nthElement(x, tf.constant(2)).asOutput();
 
-      Output<TFloat32> x1 = tf.placeholder(TFloat32.class).output();
-      Output<TFloat32> x2 = tf.placeholder(TFloat32.class).output();
-      Operand<TFloat32> x = tf.concat(Arrays.asList(x1, x2), tf.constant(0));
-      Output<TFloat32> y = tf.math.square(x).y();
-
-      Output<?>[] grads0 = g.addGradients(y, toArray(x1, x2));
+      Output<?>[] grads0 = g.addGradients(y, toArray(x));
       assertNotNull(grads0);
-      assertEquals(2, grads0.length);
+      assertEquals(1, grads0.length);
       assertEquals(DataType.DT_FLOAT, grads0[0].dataType());
-      assertEquals(DataType.DT_FLOAT, grads0[1].dataType());
 
-      try (TFloat32 c1 = TFloat32.scalarOf(3.0f);
-          TFloat32 c2 = TFloat32.scalarOf(2.0f);
+      try (TFloat32 c1 = TFloat32.vectorOf(3.0f, 2.0f, 1.0f, 0.0f);
           AutoCloseableList<Tensor> outputs =
-              new AutoCloseableList<>(
-                  s.runner().feed(x1, c1).feed(x2, c2).fetch(grads0[0]).fetch(grads0[1]).run())) {
+              new AutoCloseableList<>(s.runner().feed(x, c1).fetch(grads0[0]).run())) {
 
-        assertEquals(2, outputs.size());
-        assertEquals(6.0f, ((TFloat32) outputs.get(0)).getFloat(), 0.0f);
-        assertEquals(4.0f, ((TFloat32) outputs.get(1)).getFloat(), 0.0f);
+        assertEquals(1, outputs.size());
+        assertEquals(0.0f, ((TFloat32) outputs.get(0)).getFloat(), 0.0f);
       }
     }
   }