From edb139cd6e805173b6d3ae968437009181b34367 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 23 Feb 2021 16:28:44 -0500
Subject: [PATCH 1/5] Initial support for restoring a saved session.

---
 .../src/main/java/org/tensorflow/Session.java | 20 +++++++++++-
 .../test/java/org/tensorflow/SessionTest.java | 31 ++++++++++++++++---
 2 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
index e156491d09a..42216f31c03 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
@@ -519,7 +519,25 @@ public void runInit(){
   public void save(String prefix) {
     SaverDef saverDef = graph.saverDef();
     runner()
-        .addTarget(saverDef.getSaveTensorName())
+            .addTarget(saverDef.getSaveTensorName())
+            .feed(saverDef.getFilenameTensorName(), TString.scalarOf(prefix))
+            .run();
+  }
+
+  /**
+   * Restore the actual state of the variables of this session's graph.
+   *
+   * <p>{@code prefix} is the path where the files containing the variables state live,
+   * followed by the filename prefix. For example, if {@code prefix} is set to
+   * <i>mymodel/myvariables/variables</i>, then the files are loaded from
+   * <i>mymodel/myvariables</i> and named <i>variables.data-*-of-*</i>
+   *
+   * @param prefix prefix to restore from
+   */
+  public void restore(String prefix) {
+    SaverDef saverDef = graph.saverDef();
+    runner()
+        .addTarget(saverDef.getRestoreOpName())
         .feed(saverDef.getFilenameTensorName(), TString.scalarOf(prefix))
         .run();
   }
diff --git a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
index b1928bff51c..4572b7232fd 100644
--- a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
+++ b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
@@ -20,10 +20,14 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Comparator;
+import java.util.List;
+
 import org.junit.jupiter.api.Test;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.core.Init;
@@ -32,6 +36,7 @@
 import org.tensorflow.op.linalg.MatMul;
 import org.tensorflow.op.math.Add;
 import org.tensorflow.proto.framework.ConfigProto;
+import org.tensorflow.proto.framework.GraphDef;
 import org.tensorflow.proto.framework.RunOptions;
 import org.tensorflow.ndarray.Shape;
 import org.tensorflow.ndarray.NdArrays;
@@ -208,21 +213,39 @@ public void runInitByName() {
   }
 
   @Test
-  public void save() throws IOException  {
-    Path testFolder = Files.createTempDirectory("tf-session-save-test");
+  public void saveAndRestore() throws IOException  {
+    Path testFolder = Files.createTempDirectory("tf-session-save-restore-test");
     try (Graph g = new Graph()) {
       Ops tf = Ops.create(g);
-      Variable<TFloat32> x = tf.variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
-      Variable<TFloat32> y = tf.variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
+      Variable<TFloat32> x = tf.withName("x").variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
+      Variable<TFloat32> y = tf.withName("y").variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
       Init init = tf.init();
+      GraphDef graphDef = g.toGraphDef();
 
       try (Session s = new Session(g)) {
         s.run(init);
         s.save(testFolder.resolve("checkpoint").toString());
+        try (Graph restoredGraph = new Graph()) {
+          restoredGraph.importGraphDef(graphDef);
+          try (Session restoredSession = new Session(restoredGraph)) {
+            restoredSession.restore(testFolder.resolve("checkpoint").toString());
+            try (AutoCloseableList<Tensor> oldList = new AutoCloseableList<>(s.runner().fetch("x").fetch("y").run());
+                 AutoCloseableList<Tensor> newList = new AutoCloseableList<>(restoredSession.runner().fetch("x").fetch("y").run())){
+              assertEquals(oldList.get(0),newList.get(0));
+              assertEquals(oldList.get(1),newList.get(1));
+            }
+          }
+        }
       }
     }
     assertTrue(Files.exists(testFolder.resolve("checkpoint.index")));
     assertTrue(Files.exists(testFolder.resolve("checkpoint.data-00000-of-00001")));
+
+    // Cleanup test dir
+    Files.walk(testFolder)
+            .sorted(Comparator.reverseOrder())
+            .map(Path::toFile)
+            .forEach(File::delete);
   }
 
   private static RunOptions fullTraceRunOptions() {

From 03ee4e59df296e5f55a62a377496c18c1bb78d7e Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 23 Feb 2021 17:47:36 -0500
Subject: [PATCH 2/5] Updating surefire version to 3.0.0-M5 to get around a
 stdout test issue.

---
 tensorflow-core/tensorflow-core-api/pom.xml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow-core/tensorflow-core-api/pom.xml b/tensorflow-core/tensorflow-core-api/pom.xml
index f6a2ca2a14c..fd3605d5759 100644
--- a/tensorflow-core/tensorflow-core-api/pom.xml
+++ b/tensorflow-core/tensorflow-core-api/pom.xml
@@ -373,7 +373,7 @@
       </plugin>
       <plugin>
         <artifactId>maven-surefire-plugin</artifactId>
-        <version>2.22.0</version>
+        <version>3.0.0-M5</version>
         <executions>
           <execution>
             <!--
@@ -389,6 +389,8 @@
           </execution>
         </executions>
         <configuration>
+          <!-- Activate the use of TCP to transmit events to the plugin -->
+          <forkNode implementation="org.apache.maven.plugin.surefire.extensions.SurefireForkNodeFactory"/>
           <additionalClasspathElements>
             <additionalClasspathElement>${project.build.directory}/${project.artifactId}-${project.version}-${native.classifier}.jar</additionalClasspathElement>
             <!-- Note: the following path is not accessible in deploying profile, so other libraries like

From 9f50bf199c4235ece5ef41c0050794464dfe115f Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 23 Feb 2021 21:41:09 -0500
Subject: [PATCH 3/5] Updates to ensure the saverdef doesn't mutate the graph
 twice.

---
 .../src/main/java/org/tensorflow/Graph.java   | 35 +++++++++++++++----
 .../src/main/java/org/tensorflow/Session.java | 15 ++++----
 .../test/java/org/tensorflow/SessionTest.java | 11 ++++--
 3 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
index 988683895c4..82c9272e2f2 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
@@ -28,6 +28,7 @@
 
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -47,6 +48,7 @@
 import org.tensorflow.op.Op;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.Identity;
 import org.tensorflow.op.core.NoOp;
 import org.tensorflow.op.core.Placeholder;
 import org.tensorflow.op.train.Restore;
@@ -439,15 +441,32 @@ public Output<?>[] whileLoop(
    * Return the {@link SaverDef} instance used to save the state of all variables present in
    * this graph.
    *
-   * <p/>On the first call of this method, all nodes necessary to save and restore the state of the
-   * variables are added to the graph. Consequently, any variables that are added to the graph after
-   * this call could not be saved nor restored using this {@link SaverDef}.
+   * <p/> The first time this method is called it builds the {@link SaverDef}. If this graph already
+   * contains a "save/restore_all" operation then it is assumed to contain all necessary saving and
+   * restoring operations. If that operation does not exist then the graph is mutated to add all
+   * the nodes necessary to save and restore the state of the graph. Consequently, any variables
+   * that are added to the graph after this call will not be saved nor restored using this
+   * {@link SaverDef}.
    *
    * @return a {@link SaverDef} instance
    */
   synchronized SaverDef saverDef() {
     if (saverDef == null) {
-      saverDef = addVariableSaver(this);
+      // Check to see if this graph has a restore operation
+      if (operation("save/restore_all") == null) {
+        // No saver, create one by mutating the graph
+        saverDef = addVariableSaver(this);
+      } else {
+        // This graph already has saving/restoring operations,
+        // regenerate SaverDef without mutating. The names mirror
+        // the python implementation for compatibility.
+        // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/saver.py
+        saverDef = SaverDef.newBuilder()
+                           .setFilenameTensorName("save/filename")
+                           .setSaveTensorName("save/Save")
+                           .setRestoreOpName("save/restore_all")
+                           .build();
+      }
     }
     return saverDef;
   }
@@ -798,13 +817,15 @@ private static SaverDef addVariableSaver(Graph graph) {
     Constant<TString> varNamesTensor = tf.constant(StdArrays.ndCopyOf(varNames.toArray(tmp)));
     Operand<TString> varSlices = tf.zerosLike(varNamesTensor);
 
-    Placeholder<TString> saveFilename = tf.placeholder(TString.class);
+    Placeholder<TString> saveFilename = tf.withName("filename").placeholder(TString.class);
     Save saveVariables = tf.train.save(
         saveFilename,
         varNamesTensor,
         varSlices,
         varOutputs
     );
+    Identity<TString> id = tf.withSubScope("control_dependency")
+            .withControlDependencies(Arrays.asList(saveFilename,saveVariables)).identity(saveFilename);
     Restore restoreVariables = tf.train.restore(
         saveFilename,
         varNamesTensor,
@@ -815,11 +836,11 @@ private static SaverDef addVariableSaver(Graph graph) {
     for (int i = 0; i < varOutputs.size(); ++i) {
       restoreOps.add(tf.assign(varOutputs.get(i), (Operand) restoreVariables.tensors().get(i)));
     }
-    NoOp restoreAll = tf.withControlDependencies(restoreOps).noOp();
+    NoOp restoreAll = tf.withControlDependencies(restoreOps).withName("restore_all").noOp();
 
     return SaverDef.newBuilder()
         .setFilenameTensorName(saveFilename.op().name())
-        .setSaveTensorName(saveVariables.op().name())
+        .setSaveTensorName(id.op().name())
         .setRestoreOpName(restoreAll.op().name())
         .build();
   }
diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
index 42216f31c03..b67f4a611e6 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java
@@ -512,14 +512,13 @@ public void runInit(){
    * <i>mymodel/myvariables</i> and named <i>variables.data-*-of-*</i>
    *
    * <p>Note that this method might alter the underlying graph if it is the first time that one
-   * of its session is saved, see {@link Graph#saverDef()} for more details.
+   * of its sessions is saved, see {@link Graph#saverDef()} for more details.
    *
    * @param prefix prefix to the variable files to save
    */
   public void save(String prefix) {
     SaverDef saverDef = graph.saverDef();
-    runner()
-            .addTarget(saverDef.getSaveTensorName())
+    runner().addTarget(saverDef.getSaveTensorName())
             .feed(saverDef.getFilenameTensorName(), TString.scalarOf(prefix))
             .run();
   }
@@ -532,14 +531,16 @@ public void save(String prefix) {
    * <i>mymodel/myvariables/variables</i>, then the files are loaded from
    * <i>mymodel/myvariables</i> and named <i>variables.data-*-of-*</i>
    *
+   * <p>Note that this method might alter the underlying graph if it is the first time that one
+   * of its sessions is saved, see {@link Graph#saverDef()} for more details.
+   *
    * @param prefix prefix to restore from
    */
   public void restore(String prefix) {
     SaverDef saverDef = graph.saverDef();
-    runner()
-        .addTarget(saverDef.getRestoreOpName())
-        .feed(saverDef.getFilenameTensorName(), TString.scalarOf(prefix))
-        .run();
+    runner().addTarget(saverDef.getRestoreOpName())
+            .feed(saverDef.getFilenameTensorName(), TString.scalarOf(prefix))
+            .run();
   }
 
   /**
diff --git a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
index 4572b7232fd..6a65166e85a 100644
--- a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
+++ b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
@@ -20,13 +20,13 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.BufferedOutputStream;
 import java.io.File;
+import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.util.Comparator;
-import java.util.List;
 
 import org.junit.jupiter.api.Test;
 import org.tensorflow.op.Ops;
@@ -220,11 +220,16 @@ public void saveAndRestore() throws IOException  {
       Variable<TFloat32> x = tf.withName("x").variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
       Variable<TFloat32> y = tf.withName("y").variable(tf.random.randomUniform(tf.constant(Shape.of(3, 3L)), TFloat32.class));
       Init init = tf.init();
-      GraphDef graphDef = g.toGraphDef();
 
       try (Session s = new Session(g)) {
         s.run(init);
         s.save(testFolder.resolve("checkpoint").toString());
+        GraphDef graphDef = g.toGraphDef();
+
+        try (BufferedOutputStream os = new BufferedOutputStream(new FileOutputStream(testFolder.resolve("graph.pb").toFile()))) {
+          graphDef.writeTo(os);
+        }
+
         try (Graph restoredGraph = new Graph()) {
           restoredGraph.importGraphDef(graphDef);
           try (Session restoredSession = new Session(restoredGraph)) {

From 2490c80c3695e9ac82da3ccd3d337b3b89606ea9 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 23 Feb 2021 22:05:40 -0500
Subject: [PATCH 4/5] Matching the filename placeholder in saverdef.

---
 .../src/main/java/org/tensorflow/Graph.java        | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
index 82c9272e2f2..9a93ac8c6c1 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
@@ -44,6 +44,7 @@
 import org.tensorflow.internal.c_api.TF_Output;
 import org.tensorflow.internal.c_api.TF_Status;
 import org.tensorflow.internal.c_api.TF_WhileParams;
+import org.tensorflow.ndarray.Shape;
 import org.tensorflow.ndarray.StdArrays;
 import org.tensorflow.op.Op;
 import org.tensorflow.op.Ops;
@@ -51,6 +52,7 @@
 import org.tensorflow.op.core.Identity;
 import org.tensorflow.op.core.NoOp;
 import org.tensorflow.op.core.Placeholder;
+import org.tensorflow.op.core.PlaceholderWithDefault;
 import org.tensorflow.op.train.Restore;
 import org.tensorflow.op.train.Save;
 import org.tensorflow.proto.framework.GraphDef;
@@ -462,8 +464,8 @@ synchronized SaverDef saverDef() {
         // the python implementation for compatibility.
         // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/saver.py
         saverDef = SaverDef.newBuilder()
-                           .setFilenameTensorName("save/filename")
-                           .setSaveTensorName("save/Save")
+                           .setFilenameTensorName("save/Const")
+                           .setSaveTensorName("save/control_dependency")
                            .setRestoreOpName("save/restore_all")
                            .build();
       }
@@ -818,14 +820,16 @@ private static SaverDef addVariableSaver(Graph graph) {
     Operand<TString> varSlices = tf.zerosLike(varNamesTensor);
 
     Placeholder<TString> saveFilename = tf.withName("filename").placeholder(TString.class);
+    // Added for compatibility with TF 1.x
+    PlaceholderWithDefault<TString> constName = tf.withName("Const").placeholderWithDefault(saveFilename, Shape.scalar());
     Save saveVariables = tf.train.save(
         saveFilename,
         varNamesTensor,
         varSlices,
         varOutputs
     );
-    Identity<TString> id = tf.withSubScope("control_dependency")
-            .withControlDependencies(Arrays.asList(saveFilename,saveVariables)).identity(saveFilename);
+    Identity<TString> id = tf.withControlDependencies(Arrays.asList(saveFilename,saveVariables))
+            .withName("control_dependency").identity(saveFilename);
     Restore restoreVariables = tf.train.restore(
         saveFilename,
         varNamesTensor,
@@ -839,7 +843,7 @@ private static SaverDef addVariableSaver(Graph graph) {
     NoOp restoreAll = tf.withControlDependencies(restoreOps).withName("restore_all").noOp();
 
     return SaverDef.newBuilder()
-        .setFilenameTensorName(saveFilename.op().name())
+        .setFilenameTensorName(constName.op().name())
         .setSaveTensorName(id.op().name())
         .setRestoreOpName(restoreAll.op().name())
         .build();

From 230c11472efd63cc93b99218db8bfea34723b915 Mon Sep 17 00:00:00 2001
From: Adam Pocock <adam.pocock@oracle.com>
Date: Tue, 23 Feb 2021 22:24:59 -0500
Subject: [PATCH 5/5] Unpicking the nested placeholders, it makes various
 savers throw exceptions.

---
 .../src/main/java/org/tensorflow/Graph.java               | 8 ++------
 .../src/test/java/org/tensorflow/SessionTest.java         | 4 ----
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
index 9a93ac8c6c1..0f7a291466c 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
@@ -44,7 +44,6 @@
 import org.tensorflow.internal.c_api.TF_Output;
 import org.tensorflow.internal.c_api.TF_Status;
 import org.tensorflow.internal.c_api.TF_WhileParams;
-import org.tensorflow.ndarray.Shape;
 import org.tensorflow.ndarray.StdArrays;
 import org.tensorflow.op.Op;
 import org.tensorflow.op.Ops;
@@ -52,7 +51,6 @@
 import org.tensorflow.op.core.Identity;
 import org.tensorflow.op.core.NoOp;
 import org.tensorflow.op.core.Placeholder;
-import org.tensorflow.op.core.PlaceholderWithDefault;
 import org.tensorflow.op.train.Restore;
 import org.tensorflow.op.train.Save;
 import org.tensorflow.proto.framework.GraphDef;
@@ -464,7 +462,7 @@ synchronized SaverDef saverDef() {
         // the python implementation for compatibility.
         // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/saver.py
         saverDef = SaverDef.newBuilder()
-                           .setFilenameTensorName("save/Const")
+                           .setFilenameTensorName("save/filename")
                            .setSaveTensorName("save/control_dependency")
                            .setRestoreOpName("save/restore_all")
                            .build();
@@ -820,8 +818,6 @@ private static SaverDef addVariableSaver(Graph graph) {
     Operand<TString> varSlices = tf.zerosLike(varNamesTensor);
 
     Placeholder<TString> saveFilename = tf.withName("filename").placeholder(TString.class);
-    // Added for compatibility with TF 1.x
-    PlaceholderWithDefault<TString> constName = tf.withName("Const").placeholderWithDefault(saveFilename, Shape.scalar());
     Save saveVariables = tf.train.save(
         saveFilename,
         varNamesTensor,
@@ -843,7 +839,7 @@ private static SaverDef addVariableSaver(Graph graph) {
     NoOp restoreAll = tf.withControlDependencies(restoreOps).withName("restore_all").noOp();
 
     return SaverDef.newBuilder()
-        .setFilenameTensorName(constName.op().name())
+        .setFilenameTensorName(saveFilename.op().name())
         .setSaveTensorName(id.op().name())
         .setRestoreOpName(restoreAll.op().name())
         .build();
diff --git a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
index 6a65166e85a..d7ea381d315 100644
--- a/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
+++ b/tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SessionTest.java
@@ -226,10 +226,6 @@ public void saveAndRestore() throws IOException  {
         s.save(testFolder.resolve("checkpoint").toString());
         GraphDef graphDef = g.toGraphDef();
 
-        try (BufferedOutputStream os = new BufferedOutputStream(new FileOutputStream(testFolder.resolve("graph.pb").toFile()))) {
-          graphDef.writeTo(os);
-        }
-
         try (Graph restoredGraph = new Graph()) {
           restoredGraph.importGraphDef(graphDef);
           try (Session restoredSession = new Session(restoredGraph)) {