Add TensorScope

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

@@ -84,8 +84,9 @@ public String toString() {
    *
    * <p>This is only supported in an eager execution environment.
    *
+   * @param scope the {@link TensorScope} to create the tensor in
    * @param outputIdx index of the output of this operation
    * @return output tensor
    */
-  abstract Tensor tensor(int outputIdx);
+  abstract Tensor tensor(TensorScope scope, int outputIdx);
 }

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

@@ -16,9 +16,9 @@
 package org.tensorflow;
 
 import java.io.IOException;
+import java.util.HashMap;
 import java.util.List;
 import java.util.ListIterator;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 import org.tensorflow.op.Ops;
@@ -43,12 +43,12 @@ public class ConcreteFunction implements AutoCloseable {
    * Creates a function by building a new graph.
    *
    * <p>The {@code functionBuilder} must initialize the function graph from the provided
-   * {@link Ops} instance and return a valid signature that will be used to feed the input tensors
-   * and fetch the output tensors on execution.
+   * {@link Ops} instance and return a valid signature that will be used to feed the input tensors and fetch the output
+   * tensors on execution.
    *
    * <p>The function will be the owner of the new graph and its resulting session. Therefore,
-   * the function must be enclosed properly with a try-with-resources block to guarantee that
-   * all native resources will be freed once the function is discarded. For example:
+   * the function must be enclosed properly with a try-with-resources block to guarantee that all native resources will
+   * be freed once the function is discarded. For example:
    *
    * <pre>{@code
    * public class MyModel {
@@ -87,8 +87,8 @@ public static ConcreteFunction create(Function<Ops, Signature> functionBuilder)
    * Create a function from a signature and an existing graph.
    *
    * <p>The function will keep the ownership of the session used to run the graph but not
-   * the graph itself, meaning that the lifetime of the latter can extend beyond the scope
-   * of the function. For example:
+   * the graph itself, meaning that the lifetime of the latter can extend beyond the scope of the function. For
+   * example:
    *
    * <pre>{@code
    * try (Graph g = new Graph()) {
@@ -116,8 +116,8 @@ public static ConcreteFunction create(Signature signature, Graph graph) {
    * Create a function from a signature and a valid graph session.
    *
    * <p>The function will not own the session nor its graph, meaning that their lifetime
-   * can extend beyond the scope of the function. Therefore the function does not need to be
-   * closed after its usage. For example:
+   * can extend beyond the scope of the function. Therefore the function does not need to be closed after its usage. For
+   * example:
    *
    * <pre>{@code
    * try (Graph g = new Graph()) {
@@ -156,14 +156,11 @@ public Signature signature() {
   /**
    * Invokes a function.
    *
-   * <p>Caller is responsible for closing all Tensors.
-   *
-   * @param arguments list of tensors to pass in input to the function,
-   *                  mapped by their signature name
-   * @return output tensors resulting from the execution of the function,
-   *         mapped by their signature name
+   * @param scope the {@link TensorScope} to create the outputs in
+   * @param arguments list of tensors to pass in input to the function, mapped by their signature name
+   * @return output tensors resulting from the execution of the function, mapped by their signature name
    */
-  public Map<String, Tensor> call(Map<String, Tensor> arguments)
+  public Map<String, Tensor> call(TensorScope scope, Map<String, Tensor> arguments)
       throws IllegalArgumentException {
 
     final SignatureDef signatureDef = signature.asSignatureDef();
@@ -180,13 +177,13 @@ public Map<String, Tensor> call(Map<String, Tensor> arguments)
     Map<String, TensorInfo> outputToNode = signatureDef.getOutputsMap();
     outputToNode.values().forEach(t -> runner.fetch(t.getName()));
 
-    List<Tensor> resultTensors = runner.run();
+    List<Tensor> resultTensors = runner.run(scope);
     try {
       ListIterator<Tensor> resultTensorIter = resultTensors.listIterator();
       Map<String, Tensor> returnMap = new HashMap<String, Tensor>();
 
       // Use the output names as present in the signature definition
-      for (String nodeName: outputToNode.keySet()) {
+      for (String nodeName : outputToNode.keySet()) {
         returnMap.put(nodeName, resultTensorIter.next());
       }
       return returnMap;
@@ -203,29 +200,27 @@ public Map<String, Tensor> call(Map<String, Tensor> arguments)
   /**
    * Invokes a function with a single input and output.
    *
-   * <p>Caller is responsible for closing all Tensors.
-   *
+   * @param scope the {@link TensorScope} to create the output in
    * @param tensor input tensor
    * @return output tensor
-   * @throws IllegalArgumentException if there are multiple input or output parameters defined
-   *                                  in the function
+   * @throws IllegalArgumentException if there are multiple input or output parameters defined in the function
    */
-  public Tensor call(Tensor tensor) throws IllegalArgumentException {
+  public Tensor call(TensorScope scope, Tensor tensor) throws IllegalArgumentException {
     final SignatureDef signatureDef = signature.asSignatureDef();
 
     if (signatureDef.getInputsCount() != 1) {
       throw new IllegalArgumentException(
-        String.format("Function [%s] requires multiple inputs", signatureDef.getMethodName()));
+          String.format("Function [%s] requires multiple inputs", signatureDef.getMethodName()));
     }
     String inputNodeName = signatureDef.getInputsMap().values().iterator().next().getName();
 
     if (signatureDef.getOutputsCount() != 1) {
       throw new IllegalArgumentException(
-        String.format("Function [%s] has multiple outputs", signatureDef.getMethodName()));
+          String.format("Function [%s] has multiple outputs", signatureDef.getMethodName()));
     }
     String outputNodeName = signatureDef.getOutputsMap().values().iterator().next().getName();
 
-    return session.runner().feed(inputNodeName, tensor).fetch(outputNodeName).run().get(0);
+    return session.runner().feed(inputNodeName, tensor).fetch(outputNodeName).run(scope).get(0);
   }
 
   /**
@@ -245,8 +240,8 @@ public void save(String exportDir) throws IOException {
    * Returns the session used to execute the graph when calling this function
    *
    * <p>In general, a user does not need to handle directly the session of a function and rely
-   * on {@link #call(Map)} to execute the graph instead. But in some cases, direct access to
-   * the session might be necessary, as it allows more running options.
+   * on {@link #call(TensorScope, Map)} to execute the graph instead. But in some cases, direct access to the session
+   * might be necessary, as it allows more running options.
    *
    * @return the function session
    */

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

@@ -35,8 +35,8 @@
  * Implementation of an {@link Operation} executed eagerly.
  *
  * <p>EagerOperation instances are valid only as long as the {@link EagerSession} they are a part of
- * is valid. Thus, if {@link EagerSession#close()} has been invoked, then methods on the
- * EagerOperation instance may fail with an {@code IllegalStateException}.
+ * is valid. Thus, if {@link EagerSession#close()} has been invoked, then methods on the EagerOperation instance may
+ * fail with an {@code IllegalStateException}.
  *
  * <p>EagerOperation instances are thread-safe.
  */
@@ -120,10 +120,10 @@ DataType dtype(int outputIndex) {
   }
 
   @Override
-  Tensor tensor(int outputIndex) {
+  Tensor tensor(TensorScope scope, int outputIndex) {
     Tensor tensor = outputTensors.get(outputIndex);
     if (tensor == null) {
-      tensor = resolveTensor(outputIndex);
+      tensor = resolveTensor(scope, outputIndex);
     }
     return tensor;
   }
@@ -133,11 +133,11 @@ Tensor tensor(int outputIndex) {
   private final String name;
   private final AtomicReferenceArray<Tensor> outputTensors;
 
-  private Tensor resolveTensor(int outputIndex) {
+  private Tensor resolveTensor(TensorScope scope, int outputIndex) {
     // Take an optimistic approach, where we attempt to resolve the output tensor without locking.
     // If another thread has resolved it meanwhile, release our copy and reuse the existing one
     // instead.
-    Tensor tensor = resolveTensorHandle(getUnsafeNativeHandle(outputIndex), session);
+    Tensor tensor = resolveTensorHandle(getUnsafeNativeHandle(outputIndex), scope);
     if (!outputTensors.compareAndSet(outputIndex, null, tensor)) {
       session.detach(tensor.asRawTensor().nativeHandle());
       tensor = outputTensors.get(outputIndex);
@@ -160,13 +160,13 @@ private static void requireTensorHandle(TFE_TensorHandle handle) {
     }
   }
 
-  private static Tensor resolveTensorHandle(TFE_TensorHandle handle, EagerSession session) {
+  private static Tensor resolveTensorHandle(TFE_TensorHandle handle, TensorScope tensorScope) {
     requireTensorHandle(handle);
     try (PointerScope scope = new PointerScope()) {
       TF_Status status = TF_Status.newStatus();
-      TF_Tensor tensor = TFE_TensorHandleResolve(handle, status).withDeallocator();
+      TF_Tensor tensor = TFE_TensorHandleResolve(handle, status).withDeallocator(true);
       status.throwExceptionIfNotOK();
-      return RawTensor.fromHandle(tensor, session).asTypedTensor();
+      return RawTensor.fromHandle(tensorScope, tensor).asTypedTensor();
     }
   }
 

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

@@ -37,8 +37,8 @@
  * Implementation for an {@link Operation} added as a node to a {@link Graph}.
  *
  * <p>GraphOperation instances are valid only as long as the {@link Graph} they are a part of is
- * valid. Thus, if {@link Graph#close()} has been invoked, then methods on the GraphOperation
- * instance may fail with an {@code IllegalStateException}.
+ * valid. Thus, if {@link Graph#close()} has been invoked, then methods on the GraphOperation instance may fail with an
+ * {@code IllegalStateException}.
  *
  * <p>GraphOperation instances are immutable and thread-safe.
  */
@@ -166,7 +166,7 @@ DataType dtype(int outputIdx) {
   }
 
   @Override
-  Tensor tensor(int outputIdx) {
+  Tensor tensor(TensorScope scope, int outputIdx) {
     throw new IllegalStateException("Graph tensors must be fetched by running a session");
   }
 
@@ -236,7 +236,9 @@ private static long[] shape(TF_Graph graphHandle, TF_Operation opHandle, int out
       TF_Status status = TF_Status.newStatus();
       int numDims = TF_GraphGetTensorNumDims(graphHandle, output, status);
       status.throwExceptionIfNotOK();
-      if (numDims < 0) return null;
+      if (numDims < 0) {
+        return null;
+      }
       long[] dims = new long[numDims];
       TF_GraphGetTensorShape(graphHandle, output, dims, numDims, status);
       status.throwExceptionIfNotOK();
@@ -250,8 +252,8 @@ private static int dtype(TF_Graph graphHandle, TF_Operation opHandle, int output
 
     int numOutputs = TF_OperationNumOutputs(opHandle);
     if (outputIndex < 0 || outputIndex >= numOutputs) {
-        throw new IndexOutOfBoundsException("invalid output index (" + outputIndex
-            + ") for an operation that has " + numOutputs + " outputs");
+      throw new IndexOutOfBoundsException("invalid output index (" + outputIndex
+          + ") for an operation that has " + numOutputs + " outputs");
     }
 
     try (PointerScope scope = new PointerScope()) {

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

@@ -58,11 +58,12 @@ public interface Operand<T extends TType> extends Op, Shaped {
    *
    * <i>Only works when running in an eager execution</i>
    *
+   * @param scope the {@link TensorScope} to create the tensor in
    * @return the tensor
    * @throws IllegalStateException if this is an operand of a graph
    */
-  default T asTensor() {
-    return asOutput().asTensor();
+  default T asTensor(TensorScope scope) {
+    return asOutput().asTensor(scope);
   }
 
   /**

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

@@ -33,31 +33,36 @@
  */
 public final class Output<T extends TType> implements Operand<T> {
 
-  /** Returns the index into the outputs of the Operation. */
+  /**
+   * Returns the index into the outputs of the Operation.
+   */
   public int index() {
     return index;
   }
 
-  /** Returns the DataType of the tensor referred to by this Output. */
+  /**
+   * Returns the DataType of the tensor referred to by this Output.
+   */
   @SuppressWarnings("unchecked")
   public DataType dataType() {
     return operation.dtype(index);
   }
 
-  /** Returns the type of the tensor referred to by this Output. */
+  /**
+   * Returns the type of the tensor referred to by this Output.
+   */
   @SuppressWarnings("unchecked")
   @Override
   public Class<T> type() {
-    return (Class<T>)TensorTypeRegistry.find(dataType()).type();
+    return (Class<T>) TensorTypeRegistry.find(dataType()).type();
   }
 
   /**
-   * Returns this Output object with the type {@code Output<U>}. This method is useful when given a
-   * value of type {@code Output<?>}.
+   * Returns this Output object with the type {@code Output<U>}. This method is useful when given a value of type {@code
+   * Output<?>}.
    *
    * @param type any supported tensor type
-   * @throws IllegalArgumentException if the actual data type of this object does not match the type
-   *     {@code U}.
+   * @throws IllegalArgumentException if the actual data type of this object does not match the type {@code U}.
    */
   @SuppressWarnings("unchecked")
   public <U extends TType> Output<U> expect(Class<U> type) {
@@ -72,8 +77,7 @@ public <U extends TType> Output<U> expect(Class<U> type) {
    * Returns the tensor at this output.
    *
    * <p>This operation is only supported on the outputs of an operation executed eagerly. For graph
-   * environments, output tensors must be fetched by running a session, using {@link
-   * Session.Runner#fetch(Output)}.
+   * environments, output tensors must be fetched by running a session, using {@link Session.Runner#fetch(Output)}.
    *
    * <p>It is recommended to close explicitly the returned tensor as soon as possible, since the
    * garbage collector is not aware of the amount of memory it consumes, which can be significant.
@@ -84,8 +88,8 @@ public <U extends TType> Output<U> expect(Class<U> type) {
    * @see EagerSession
    */
   @SuppressWarnings("unchecked")
-  public T asTensor() {
-    return (T)operation.tensor(index);
+  public T asTensor(TensorScope scope) {
+    return (T) operation.tensor(scope, index);
   }
 
   /**
@@ -130,7 +134,9 @@ public String toString() {
         operation.type(), operation.name(), index, shape().toString(), dataType());
   }
 
-  /** Handle to the idx-th output of the Operation {@code op}. */
+  /**
+   * Handle to the idx-th output of the Operation {@code op}.
+   */
   Output(AbstractOperation op, int idx) {
     operation = op;
     index = idx;