Working gradients

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

Author: rnett

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

@@ -122,6 +122,7 @@ public void close() {
       }
       delete(nativeHandle);
       nativeHandle = null;
+      allGraphs.remove(this);
     }
   }
 
@@ -1320,7 +1321,7 @@ private static SaverDef addVariableSaver(Graph graph) {
    */
   public static Graph findGraphForPointer(NativeGraphPointer pointer) {
     for (Graph g : allGraphs) {
-      if (g.nativeHandle.graph().equals(pointer)) {
+      if (g.nativeHandle != null && !g.nativeHandle.isNull() && g.nativeHandle.graph().equals(pointer)) {
         return g;
       }
     }

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

@@ -22,6 +22,7 @@
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_FinishOperationLocked;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewOperation;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewOperationLocked;
+import static org.tensorflow.internal.c_api.global.tensorflow.TF_OperationName;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrBool;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrBoolList;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrFloat;
@@ -457,7 +458,7 @@ private static TF_Operation finishDangerousGradient(TF_Graph g, TF_OperationDesc
       TF_Status status = TF_Status.newStatus();
       TF_Operation op = TF_FinishOperationLocked(handle, status);
       status.throwExceptionIfNotOK();
-//      g.name_map().put(TF_OperationName(op), null);
+      g.name_map().erase(TF_OperationName(op));
       return op;
     }
   }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/presets/tensorflow.java

@@ -287,7 +287,7 @@ public void init(ClassProperties properties) {
         .put(new Info("TF_Graph::refiner", "TF_Graph::mu",
             "TF_Graph::sessions", "TF_Graph::delete_requested").skip())
         .put(new Info("std::unordered_map<tensorflow::string,tensorflow::Node*>")
-            .pointerTypes("NameMap").define())
+            .pointerTypes("NameMap").define().javaText("public native long erase(@StdString BytePointer key);"))
         .put(new Info("TF_Function")
           .pointerTypes("TF_Function")
           .base("org.tensorflow.internal.c_api.AbstractTF_Function"))

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

@@ -38,7 +38,7 @@ public void addGradientsToGraph() {
             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);
+            return Arrays.asList(a, b, tf.constant(0f));
           });
       Ops tf = Ops.create(g);
 
@@ -64,8 +64,8 @@ public void addGradientsToGraph() {
                   .run())) {
 
         assertEquals(2, outputs.size());
-        assertEquals(3.0f, ((TFloat32) outputs.get(0)).getFloat(), 0.0f);
-        assertEquals(2.0f, ((TFloat32) outputs.get(1)).getFloat(), 0.0f);
+        assertEquals(6.0f, ((TFloat32) outputs.get(0)).getFloat(), 0.0f);
+        assertEquals(4.0f, ((TFloat32) outputs.get(1)).getFloat(), 0.0f);
       }
     }
   }