Throw exception on custom gradient registration on Windows

Author: karllessard

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

@@ -26,6 +26,7 @@
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.Collections;
 import java.util.IdentityHashMap;
+import java.util.Locale;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.bytedeco.javacpp.PointerPointer;
@@ -193,9 +194,8 @@ private static synchronized boolean hasGradient(String opType) {
    * <p>Note that this only works with graph gradients, and will eventually be deprecated in favor
    * of unified gradient support once it is fully supported by tensorflow core.
    *
-   * <p><i>Warning: Custom gradient registration is currently not supported on Windows and may crash
-   * the JVM if trying to do so. See <a href=https://github.com/tensorflow/java/issues/486>GitHub
-   * issue</a> related to this.</i>
+   * <p><i>Warning: Custom gradient registration is currently not supported on Windows, see <a
+   * href=https://github.com/tensorflow/java/issues/486>GitHub issue</a> for more info.</i>
    *
    * @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}.
@@ -205,6 +205,10 @@ private static synchronized boolean hasGradient(String opType) {
    */
   public static synchronized boolean registerCustomGradient(
       String opType, RawCustomGradient gradient) {
+    if (isWindowsOs()) {
+      throw new UnsupportedOperationException(
+          "Custom gradient registration is not supported on Windows systems.");
+    }
     if (hasGradient(opType)) {
       return false;
     }
@@ -220,9 +224,8 @@ public static synchronized boolean registerCustomGradient(
    * generated op classes or custom op classes with the correct annotations. To operate on the
    * {@link org.tensorflow.GraphOperation} directly use {@link RawCustomGradient}.
    *
-   * <p><i>Warning: Custom gradient registration is currently not supported on Windows and may crash
-   * the JVM if trying to do so. See <a href=https://github.com/tensorflow/java/issues/486>GitHub
-   * issue</a> related to this.</i>
+   * <p><i>Warning: Custom gradient registration is currently not supported on Windows, see <a
+   * href=https://github.com/tensorflow/java/issues/486>GitHub issue</a> for more info.</i>
    *
    * @param inputClass the inputs class of op to register the gradient for.
    * @param gradient the gradient function to use
@@ -233,8 +236,11 @@ public static synchronized boolean registerCustomGradient(
    */
   public static synchronized <T extends RawOpInputs<?>> boolean registerCustomGradient(
       Class<T> inputClass, CustomGradient<T> gradient) {
+    if (isWindowsOs()) {
+      throw new UnsupportedOperationException(
+          "Custom gradient registration is not supported on Windows systems.");
+    }
     OpInputsMetadata metadata = inputClass.getAnnotation(OpInputsMetadata.class);
-
     if (metadata == null) {
       throw new IllegalArgumentException(
           "Inputs Class "
@@ -261,4 +267,8 @@ public static synchronized <T extends RawOpInputs<?>> boolean registerCustomGrad
     gradientFuncs.add(g);
     return true;
   }
+
+  private static boolean isWindowsOs() {
+    return System.getProperty("os.name", "").toLowerCase(Locale.ENGLISH).startsWith("win");
+  }
 }

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

@@ -21,6 +21,7 @@
 import java.util.Arrays;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.DisabledOnOs;
+import org.junit.jupiter.api.condition.EnabledOnOs;
 import org.junit.jupiter.api.condition.OS;
 import org.tensorflow.ndarray.index.Indices;
 import org.tensorflow.op.Ops;
@@ -44,8 +45,8 @@ public void testAlreadyExisting() {
             }));
   }
 
-  // FIXME: Since TF 2.10.1, this test is failing on Windows, because the whole JVM crashes when
-  // calling the JavaCPP generated binding `NameMap.erase`. Disable it until we find a fix.
+  // FIXME: Since TF 2.10.1, custom gradient registration is failing on Windows, see
+  //        https://github.com/tensorflow/java/issues/486
   @DisabledOnOs(OS.WINDOWS)
   @Test
   public void testCustomGradient() {
@@ -76,6 +77,26 @@ public void testCustomGradient() {
     }
   }
 
+  @EnabledOnOs(OS.WINDOWS)
+  @Test
+  public void testCustomGradientThrowsOnWindows() {
+    assertThrows(
+        UnsupportedOperationException.class,
+        () ->
+            TensorFlow.registerCustomGradient(
+                NthElement.OP_NAME,
+                (tf, op, gradInputs) ->
+                    Arrays.asList(tf.withName("inAGrad").constant(0f), tf.constant(0f))));
+
+    assertThrows(
+        UnsupportedOperationException.class,
+        () ->
+            TensorFlow.registerCustomGradient(
+                NthElement.Inputs.class,
+                (tf, op, gradInputs) ->
+                    Arrays.asList(tf.withName("inAGrad").constant(0f), tf.constant(0f))));
+  }
+
   private static Output<?>[] toArray(Output<?>... outputs) {
     return outputs;
   }