Bug fixes

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

Author: rnett

tensorflow-core/tensorflow-core-api/src/gen/annotations/org/tensorflow/op/Ops.java

@@ -18,13 +18,15 @@
 package org.tensorflow.op;
 
 import java.nio.charset.Charset;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import org.tensorflow.ConcreteFunction;
 import org.tensorflow.DeviceSpec;
 import org.tensorflow.EagerSession;
 import org.tensorflow.ExecutionEnvironment;
 import org.tensorflow.Operand;
+import org.tensorflow.Operation;
 import org.tensorflow.ndarray.BooleanNdArray;
 import org.tensorflow.ndarray.ByteNdArray;
 import org.tensorflow.ndarray.DoubleNdArray;
@@ -8254,6 +8256,33 @@ public Ops withControlDependencies(Iterable<Op> controls) {
     return new Ops(scope.withControlDependencies(controls));
   }
 
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * @see {@link Scope#withControlDependencies(Iterable<Op<?>>)}
+   */
+  public Ops withControlDependencies(Op... controls) {
+    return withControlDependencies(Arrays.asList(controls));
+  }
+
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * @see {@link Scope#withControlDependencyOps(Iterable<Operation>)}
+   */
+  public Ops withControlDependencyOps(Iterable<Operation> controls) {
+    return new Ops(scope.withControlDependencyOps(controls));
+  }
+
+  /**
+   * Returns an API that adds operations to the graph with the provided control dependencies.
+   *
+   * @see {@link Scope#withControlDependencyOps(Iterable<Operation>)}
+   */
+  public Ops withControlDependencyOps(Operation... controls) {
+    return withControlDependencyOps(Arrays.asList(controls));
+  }
+
   /**
    * Returns the current {@link Scope scope} of this API
    */

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

@@ -1244,6 +1244,8 @@ private static Object[] whileLoop(
   private static SaverDef addVariableSaver(Graph graph) {
     Ops tf = Ops.create(graph).withSubScope(SAVER_DEF_SCOPE);
 
+    // TODO handle resource variables, too
+
     List<String> varNames = new ArrayList<>();
     List<Operand<?>> varOutputs = new ArrayList<>();
     List<Class<? extends TType>> varTypes = new ArrayList<>();

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

@@ -281,6 +281,7 @@ public void export() throws IOException {
       functions.forEach((k, f) -> metaGraphDef.putSignatureDef(k, f.signature().asSignatureDef()));
 
       if (!functions.containsKey(INIT_OP_SIGNATURE_KEY)) {
+
         metaGraphDef.putSignatureDef(
             INIT_OP_SIGNATURE_KEY,
             SignatureDef.newBuilder()
@@ -388,7 +389,10 @@ public Session session() {
 
   /** Return the signature of all functions available in this saved model. */
   public List<Signature> signatures() {
-    return functions.values().stream().map(f -> f.signature()).collect(Collectors.toList());
+    return functions.values().stream()
+        .map(SessionFunction::signature)
+        .filter(signature -> !signature.key().equals(INIT_OP_SIGNATURE_KEY))
+        .collect(Collectors.toList());
   }
 
   /**
@@ -593,6 +597,9 @@ private static SavedModelBundle load(
         throw new TensorFlowException("Cannot parse MetaGraphDef protocol buffer", e);
       }
     }
+    bundle.session.initialize();
+
+    //    bundle.session.restore(exportDir + "/variables/variables");
 
     return bundle;
   }

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

@@ -180,10 +180,8 @@ public void exportFunctionWithVariables() throws IOException {
       assertEquals("save/control_dependency", saverDef.getSaveTensorName());
       assertEquals("save/restore_all", saverDef.getRestoreOpName());
 
-      assertEquals(1, savedModel.metaGraphDef().getSignatureDefCount());
-      assertEquals(
-          Signature.DEFAULT_KEY,
-          savedModel.metaGraphDef().getSignatureDefMap().keySet().iterator().next());
+      assertEquals(2, savedModel.metaGraphDef().getSignatureDefCount());
+      assertTrue(savedModel.metaGraphDef().getSignatureDefMap().containsKey(Signature.DEFAULT_KEY));
 
       TensorFunction function = savedModel.function(Signature.DEFAULT_KEY);
       assertNotNull(function);

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/Names.java

@@ -1,26 +1,25 @@
 /*
-  Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+ Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-     http://www.apache.org/licenses/LICENSE-2.0
+    http://www.apache.org/licenses/LICENSE-2.0
 
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
- ==============================================================================
- */
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================
+*/
 package org.tensorflow;
 
 import com.squareup.javapoet.ArrayTypeName;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
-import java.util.Arrays;
 
 public class Names {
 
@@ -52,9 +51,12 @@ public class Names {
   public static final ClassName RawOp = ClassName.get(OpPackage, "RawOp");
   public static final ClassName Operation = ClassName.get(TensorflowPackage, "Operation");
   public static final ClassName Operands = ClassName.get(OpPackage, "Operands");
-  public static final ClassName OperationBuilder = ClassName.get(TensorflowPackage, "OperationBuilder");
-  public static final TypeName IterableOp = ParameterizedTypeName.get(ClassName.get(Iterable.class), Op);
-  public static final TypeName IterableOperation = ParameterizedTypeName.get(ClassName.get(Iterable.class), Operation);
+  public static final ClassName OperationBuilder =
+      ClassName.get(TensorflowPackage, "OperationBuilder");
+  public static final TypeName IterableOp =
+      ParameterizedTypeName.get(ClassName.get(Iterable.class), Op);
+  public static final TypeName IterableOperation =
+      ParameterizedTypeName.get(ClassName.get(Iterable.class), Operation);
   public static final TypeName ArrayOp = ArrayTypeName.of(Op);
   public static final TypeName ArrayOperation = ArrayTypeName.of(Operation);
 
@@ -63,7 +65,8 @@ public class Names {
 
   public static final ClassName Shape = ClassName.get(TensorflowPackage + ".ndarray", "Shape");
   public static final ClassName Tensor = ClassName.get(TensorflowPackage, "Tensor");
-  public static final ClassName ConcreteFunction = ClassName.get(TensorflowPackage, "ConcreteFunction");
+  public static final ClassName ConcreteFunction =
+      ClassName.get(TensorflowPackage, "ConcreteFunction");
 
   public static final ClassName Scope = ClassName.get(OpPackage, "Scope");
   public static final TypeName DeviceSpec = ClassName.get(TensorflowPackage, "DeviceSpec");
@@ -75,5 +78,4 @@ public class Names {
 
   public static final TypeName String = ClassName.get(String.class);
   public static final ClassName Arrays = ClassName.get(java.util.Arrays.class);
-
 }

tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/processor/operator/OperatorProcessor.java

@@ -625,7 +625,7 @@ private static TypeSpec buildTopClass(OpsSpec spec) {
             .addParameter(Names.ArrayOp, "controls")
             .varargs()
             .returns(Names.Ops)
-            .addStatement("return withControlDependencies(%T.asList(controls))", Names.Arrays)
+            .addStatement("return withControlDependencies($T.asList(controls))", Names.Arrays)
             .addJavadoc(
                 "Returns an API that adds operations to the graph with the provided control dependencies.\n\n"
                     + "@see {@link $T#withControlDependencies(Iterable<Op<?>>)}\n",
@@ -650,7 +650,7 @@ private static TypeSpec buildTopClass(OpsSpec spec) {
             .addParameter(Names.ArrayOperation, "controls")
             .varargs()
             .returns(Names.Ops)
-            .addStatement("return withControlDependencyOps(%T.asList(controls))", Names.Arrays)
+            .addStatement("return withControlDependencyOps($T.asList(controls))", Names.Arrays)
             .addJavadoc(
                 "Returns an API that adds operations to the graph with the provided control dependencies.\n\n"
                     + "@see {@link $T#withControlDependencyOps(Iterable<Operation>)}\n",