Clarify the difference between CustomGradient and RawCustomGradient

Signed-off-by: Ryan Nett <JNett96@gmail.com>

Author: rnett

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

@@ -152,7 +152,8 @@ private TensorFlow() {}
   }
 
   /**
-   * Keeps references to custom gradient functions to prevent them from being deallocated.
+   * Keeps references to custom gradient functions to prevent them from being deallocated. All
+   * access of this set should be synchronized on this class.
    *
    * <p><b>Required for correctness</b>
    */
@@ -170,6 +171,9 @@ private static synchronized boolean hasGradient(String opType) {
   /**
    * Register a custom gradient function for ops of {@code opType} type.
    *
+   * <p>Creates the gradient based off of a {@link GraphOperation}. To operate on the op input class
+   * instead use {@link CustomGradient}.
+   *
    * <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.
    *
@@ -193,7 +197,8 @@ public static synchronized boolean registerCustomGradient(
   /**
    * 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.
+   * generated op classes or custom op classes with the correct annotations. To operate on the
+   * {@link org.tensorflow.GraphOperation} directly use {@link RawCustomGradient}.
    *
    * @param inputClass the inputs class of op to register the gradient for.
    * @param gradient the gradient function to use

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/CustomGradient.java

@@ -26,11 +26,15 @@
  * A custom gradient for ops of type {@link T}. Should be registered using {@link
  * TensorFlow#registerCustomGradient(Class, CustomGradient)}.
  *
+ * <p>Creates the gradient based off of an instance of the op inputs class, which is created using
+ * reflection. To operate on the {@link org.tensorflow.GraphOperation} directly use {@link
+ * RawCustomGradient}.
+ *
  * @param <T> the type of op this gradient is for.
  */
 @SuppressWarnings("rawtypes")
 @FunctionalInterface
-public interface CustomGradient<T extends RawOpInputs> {
+public interface CustomGradient<T extends RawOpInputs<?>> {
 
   /**
    * Calculate the gradients for {@code op}.

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/op/RawCustomGradient.java

@@ -27,6 +27,9 @@
  * A custom gradient for an op of unspecified type. Should be registered using {@link
  * TensorFlow#registerCustomGradient(String, RawCustomGradient)}.
  *
+ * <p>Creates the gradient based off of a {@link GraphOperation}. To operate on the op input class
+ * instead use {@link CustomGradient}.
+ *
  * <p>The op type of {@code op} will depend on the op type string passed to the registration method.
  * Note that the registration method can be called more than once, resulting this gradient function
  * being used for multiple different op types.