Fixes

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

Author: rnett

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

@@ -551,8 +551,13 @@ private static SavedModelBundle fromHandle(
     // Technically, init ops should be ran first, then variable restore, but that is not possible
     // since TF_Session.loadSessionFromSavedModel does it in reverse order, so we just mark them as
     // ran.
-    if(functions.containsKey(JAVA_INIT_OP_SIGNATURE_KEY)){
-      String initOpName = functions.get(JAVA_INIT_OP_SIGNATURE_KEY).getOutputs().get(JAVA_INIT_OP_SIGNATURE_KEY).name;
+    if (functions.containsKey(JAVA_INIT_OP_SIGNATURE_KEY)) {
+      String initOpName =
+          functions
+              .get(JAVA_INIT_OP_SIGNATURE_KEY)
+              .getOutputs()
+              .get(JAVA_INIT_OP_SIGNATURE_KEY)
+              .name;
       graph.registerInitOp(graph.outputOrThrow(initOpName).op());
     }
 

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

@@ -121,7 +121,7 @@ public void exportMultipleFunctions() throws IOException {
       }
     }
     try (SavedModelBundle model = SavedModelBundle.load(testFolder.toString())) {
-      assertEquals(2, model.signatures().size());
+      assertEquals(3, model.signatures().size());
       TensorFunction f1 = model.function(Signature.DEFAULT_KEY);
       assertNotNull(f1);
       try (TFloat32 x = TFloat32.tensorOf(StdArrays.ndCopyOf(new float[] {2, 2}));
@@ -319,33 +319,6 @@ public void pythonTfFunction() {
     }
   }
 
-  @Test
-  public void exportAndLoadInitializers() throws IOException {
-    Path testFolder = Files.createTempDirectory("tf-saved-model-export-test");
-    try (Graph g = new Graph();
-        Session s = new Session(g)) {
-      Ops tf = Ops.create(g);
-      Ops init = tf.withInitScope();
-      Operand<?> handle = init.withName("variable").varHandleOp(TInt32.class, Shape.scalar());
-      init.withName("init").assignVariableOp(handle, init.constant(10));
-
-      SessionFunction f =
-          SessionFunction.create(
-              Signature.builder()
-                  .key("f")
-                  .output("out", tf.withName("read").readVariableOp(handle, TInt32.class))
-                  .build(),
-              s);
-
-      SavedModelBundle.exporter(testFolder.toString()).withFunction(f).export();
-    }
-
-    try (SavedModelBundle savedModel = SavedModelBundle.load(testFolder.toString())) {
-      TInt32 tensor = (TInt32) savedModel.session().runner().fetch("read", 0).run().get(0);
-      assertEquals(10, tensor.getInt());
-    }
-  }
-
   private static Signature buildGraphWithVariables(Ops tf, Shape xShape) {
     Placeholder<TFloat32> x = tf.placeholder(TFloat32.class, Placeholder.shape(xShape));
     Variable<TFloat32> y =