Fold name handling into ConcreteFunction, fix tests

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

Author: rnett

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

@@ -21,11 +21,10 @@
 import java.util.List;
 import java.util.Map;
 import org.tensorflow.ConcreteFunction;
+import org.tensorflow.DefinedFunction;
 import org.tensorflow.DeviceSpec;
 import org.tensorflow.EagerSession;
 import org.tensorflow.ExecutionEnvironment;
-import org.tensorflow.GraphFunction;
-import org.tensorflow.NamedGraphFunction;
 import org.tensorflow.Operand;
 import org.tensorflow.ndarray.BooleanNdArray;
 import org.tensorflow.ndarray.ByteNdArray;
@@ -1072,100 +1071,18 @@ public Bucketize bucketize(Operand<? extends TNumber> input, List<Float> boundar
   }
 
   /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present. The inputs and outputs
-   *  are keyed by the names set in the {@code ConcreteFunction}'s {@code Signature}.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public Map<String, Operand<?>> callConcreteFunction(ConcreteFunction function,
-      Operand<?>... inputs) {
-    return Function.callConcreteFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present. The inputs and outputs
-   *  are keyed by the names set in the {@code ConcreteFunction}'s {@code Signature}.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public Map<String, Operand<?>> callConcreteFunction(ConcreteFunction function,
-      List<Operand<?>> inputs) {
-    return Function.callConcreteFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present. The inputs and outputs
-   *  are keyed by the names set in the {@code ConcreteFunction}'s {@code Signature}.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public Map<String, Operand<?>> callConcreteFunction(ConcreteFunction function,
-      Map<String, Operand<?>> inputs) {
-    return Function.callConcreteFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
+   * empty
    */
-  public List<Operand<?>> callFunction(GraphFunction function, List<Operand<?>> inputs) {
-    return Function.callFunction(scope, function, inputs);
+  public Map<String, Operand<?>> call(ConcreteFunction function,
+      Map<String, Operand<?>> arguments) {
+    return Function.call(scope, function, arguments);
   }
 
   /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public List<Operand<?>> callFunction(GraphFunction function, Operand<?>... inputs) {
-    return Function.callFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public Map<String, Operand<?>> callNamedFunction(NamedGraphFunction function,
-      Map<String, Operand<?>> inputs) {
-    return Function.callNamedFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
-   */
-  public Map<String, Operand<?>> callNamedFunction(NamedGraphFunction function,
-      Operand<?>... inputs) {
-    return Function.callNamedFunction(scope, function, inputs);
-  }
-
-  /**
-   * Call {@code function}, adding it to the execution environment if it isn't already present.
-   *
-   * @param function the function to call
-   * @param inputs the inputs to the function
-   * @return the outputs of the function
+   * empty
    */
-  public Map<String, Operand<?>> callNamedFunction(NamedGraphFunction function,
-      List<Operand<?>> inputs) {
-    return Function.callNamedFunction(scope, function, inputs);
+  public List<Operand<?>> call(DefinedFunction function, List<Operand<?>> arguments) {
+    return Function.call(scope, function, arguments);
   }
 
   /**

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

@@ -39,6 +39,17 @@ public Output<?>[] outputList(int idx, int length) {
 
   @Override
   public <T extends TType> Output<T> output(int idx) {
+    if (getUnsafeNativeHandle(idx) != null && !getUnsafeNativeHandle(idx).isNull()) {
+      int numOutputs = this.numOutputs();
+      if (idx >= numOutputs) {
+        throw new IndexOutOfBoundsException(
+            "Can't get output with index " + idx + ", this op only has " + numOutputs + " outputs.");
+      }
+
+      if (idx < 0) {
+        throw new IndexOutOfBoundsException("Can't get output with index < 0.");
+      }
+    }
     return new Output<>(this, idx);
   }
 

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

@@ -17,10 +17,12 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.ListIterator;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 import org.tensorflow.op.Ops;
 import org.tensorflow.proto.framework.SignatureDef;
 import org.tensorflow.proto.framework.TensorInfo;
@@ -209,7 +211,7 @@ public Map<String, Tensor> call(Map<String, Tensor> arguments)
    * @return the outputs of the function
    */
   public Map<String, Operand<?>> call(Ops tf, Map<String, Operand<?>> arguments) {
-    return tf.callConcreteFunction(this, arguments);
+    return tf.call(this, arguments);
   }
 
   /**
@@ -255,14 +257,18 @@ public Operand<?> call(Ops tf, Operand<?> argument) {
       throw new IllegalArgumentException(
           String.format("Function [%s] requires multiple inputs", signatureDef.getMethodName()));
     }
+    String inputName = signatureDef.getInputsMap().keySet().iterator().next();
 
     if (signatureDef.getOutputsCount() != 1) {
       throw new IllegalArgumentException(
           String.format("Function [%s] has multiple outputs", signatureDef.getMethodName()));
     }
-    String outputNodeName = signatureDef.getOutputsMap().values().iterator().next().getName();
+    String outputName = signatureDef.getOutputsMap().keySet().iterator().next();
+
+    Map<String, Operand<?>> inputMap = new LinkedHashMap<>();
+    inputMap.put(inputName, argument);
 
-    return tf.callConcreteFunction(this, argument).get(outputNodeName);
+    return call(tf, inputMap).get(outputName);
   }
 
   /**
@@ -298,12 +304,29 @@ public Graph graph() {
     return graph;
   }
 
+  private DefinedFunction makeGraphFunction() {
+    String description = signature().methodName();
+    if (description == null) {
+      description = signature().key();
+    }
+
+    List<Operand<?>> inputList = signature.getInputs().values().stream()
+        .map((it) -> graph.outputOrError(it.name))
+        .collect(Collectors.toUnmodifiableList());
+
+    List<Operand<?>> outputList = signature.getOutputs().values().stream()
+        .map((it) -> graph.outputOrError(it.name))
+        .collect(Collectors.toUnmodifiableList());
+
+    return DefinedFunction.create(graph, signature.key(), true, inputList, outputList, description);
+  }
+
   /**
    * Get the graph as a function.  The graph function will be closed when this is.
    */
-  public NamedGraphFunction function() {
+  public DefinedFunction function() {
     if (function == null) {
-      function = GraphFunction.create(this, true);
+      function = makeGraphFunction();
     }
     return function;
   }
@@ -332,7 +355,7 @@ private enum Ownership {
 
   private final Graph graph;
   private final Session session;
-  private NamedGraphFunction function = null;
+  private DefinedFunction function = null;
   private final Signature signature;
   private final Ownership ownership;