Stateful/stateless processing, selector op wrappers

Signed-off-by: Ryan Nett <JNett96@gmail.com>

Author: rnett

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

@@ -222,8 +222,12 @@
 import org.tensorflow.op.core.StageClear;
 import org.tensorflow.op.core.StagePeek;
 import org.tensorflow.op.core.StageSize;
+import org.tensorflow.op.core.StatefulCase;
+import org.tensorflow.op.core.StatefulIf;
 import org.tensorflow.op.core.StatefulPartitionedCall;
+import org.tensorflow.op.core.StatefulWhile;
 import org.tensorflow.op.core.StatelessIf;
+import org.tensorflow.op.core.StatelessPartitionedCall;
 import org.tensorflow.op.core.StatelessWhile;
 import org.tensorflow.op.core.StopGradient;
 import org.tensorflow.op.core.StridedSlice;
@@ -1230,6 +1234,7 @@ public Map<String, Operand<?>> call(ConcreteFunction function,
    *  }
    *  ```
    *  </pre>
+   *  Selects between {@link StatefulCase} and {@link StatelessCase} based on the statefulness of the function arguments.
    *
    * @param branchIndex The branch selector, an int32 Tensor.
    * @param input A list of input tensors passed to the branch function.
@@ -2946,6 +2951,7 @@ public IdentityN identityN(Iterable<Operand<?>> input) {
 
   /**
    * output = cond ? then_branch(input) : else_branch(input)
+   *  Selects between {@link StatefulIf} and {@link StatelessIf} based on the statefulness of the function arguments.
    *
    * @param cond <pre>
    *    A Tensor. If the tensor is a scalar of non-boolean type, the
@@ -4015,14 +4021,16 @@ public <T extends TType> ParallelDynamicStitch<T> parallelDynamicStitch(
 
   /**
    * returns {@code f(inputs)}, where {@code f}'s body is placed and partitioned.
+   *  Selects between {@link StatefulPartitionedCall} and {@link StatelessPartitionedCall} based on the statefulness of the function arguments.
    *
    * @param args A list of input tensors.
    * @param Tout A list of output types.
    * @param f <pre>
    *    A function that takes 'args', a list of tensors, and returns 'output',
    *    another list of tensors. Input and output types are specified by 'Tin'
    *    and 'Tout'. The function body of f will be placed and partitioned across
-   *    devices, setting this op apart from the regular Call op.
+   *    devices, setting this op apart from the regular Call op. This op is
+   *    stateful.
    *  </pre>
    * @param options carries optional attribute values
    * @return a new instance of PartitionedCall
@@ -6048,6 +6056,72 @@ public StageSize stageSize(List<Class<? extends TType>> dtypes, StageSize.Option
     return StageSize.create(scope, dtypes, options);
   }
 
+  /**
+   * An n-way switch statement which calls a single branch function.
+   *  <pre>
+   *  An n-way switch statement, implementing the following:
+   *  ```
+   *  switch (branch_index) {
+   *    case 0:
+   *      output = branches[0](input);
+   *      break;
+   *    case 1:
+   *      output = branches[1](input);
+   *      break;
+   *    ...
+   *    case [[nbranches-1]]:
+   *    default:
+   *      output = branches[nbranches-1](input);
+   *      break;
+   *  }
+   *  ```
+   *  </pre>
+   *
+   * @param branchIndex The branch selector, an int32 Tensor.
+   * @param input A list of input tensors passed to the branch function.
+   * @param Tout A list of output types.
+   * @param branches <pre>
+   *    A list of functions each of which takes 'inputs' and returns a list of
+   *    tensors, whose types are the same as what every other branch returns.
+   *  </pre>
+   * @param options carries optional attribute values
+   * @return a new instance of StatefulCase
+   */
+  public StatefulCase statefulCase(Operand<TInt32> branchIndex, Iterable<Operand<?>> input,
+      List<Class<? extends TType>> Tout, List<ConcreteFunction> branches, Case.Options... options) {
+    return StatefulCase.create(scope, branchIndex, input, Tout, branches, options);
+  }
+
+  /**
+   * output = cond ? then_branch(input) : else_branch(input)
+   *
+   * @param cond <pre>
+   *    A Tensor. If the tensor is a scalar of non-boolean type, the
+   *    scalar is converted to a boolean according to the
+   *    following rule: if the scalar is a numerical value, non-zero means
+   *    `True` and zero means False; if the scalar is a string, non-empty
+   *    means `True` and empty means `False`. If the tensor is not a scalar,
+   *    being empty means False and being non-empty means True.
+   *  </pre>
+   * @param input A list of input tensors.
+   * @param Tout A list of output types.
+   * @param thenBranch <pre>
+   *    A function that takes 'inputs' and returns a list of tensors, whose
+   *    types are the same as what else_branch returns.
+   *  </pre>
+   * @param elseBranch <pre>
+   *  A function that takes 'inputs' and returns a list of tensors, whose
+   *  types are the same as what then_branch returns.
+   *  </pre>
+   * @param options carries optional attribute values
+   * @return a new instance of StatefulIf
+   */
+  public StatefulIf statefulIf(Operand<? extends TType> cond, Iterable<Operand<?>> input,
+      List<Class<? extends TType>> Tout, ConcreteFunction thenBranch, ConcreteFunction elseBranch,
+      If.Options... options) {
+    return StatefulIf.create(scope, cond, input, Tout, thenBranch, elseBranch, options);
+  }
+
   /**
    * returns {@code f(inputs)}, where {@code f}'s body is placed and partitioned.
    *
@@ -6064,11 +6138,36 @@ public StageSize stageSize(List<Class<? extends TType>> dtypes, StageSize.Option
    * @return a new instance of StatefulPartitionedCall
    */
   public StatefulPartitionedCall statefulPartitionedCall(Iterable<Operand<?>> args,
-      List<Class<? extends TType>> Tout, ConcreteFunction f,
-      StatefulPartitionedCall.Options... options) {
+      List<Class<? extends TType>> Tout, ConcreteFunction f, PartitionedCall.Options... options) {
     return StatefulPartitionedCall.create(scope, args, Tout, f, options);
   }
 
+  /**
+   * output = input; While (Cond(output)) { output = Body(output) }
+   *
+   * @param input A list of input tensors whose types are T.
+   * @param cond <pre>
+   *    A function takes 'input' and returns a tensor.  If the tensor is
+   *    a scalar of non-boolean, the scalar is converted to a boolean
+   *    according to the following rule: if the scalar is a numerical
+   *    value, non-zero means True and zero means False; if the scalar is
+   *    a string, non-empty means True and empty means False. If the
+   *    tensor is not a scalar, non-emptiness means True and False
+   *    otherwise.
+   *  </pre>
+   * @param body <pre>
+   *    A function that takes a list of tensors and returns another
+   *    list of tensors. Both lists have the same types as specified
+   *    by T.
+   *  </pre>
+   * @param options carries optional attribute values
+   * @return a new instance of StatefulWhile
+   */
+  public StatefulWhile statefulWhile(Iterable<Operand<?>> input, ConcreteFunction cond,
+      ConcreteFunction body, While.Options... options) {
+    return StatefulWhile.create(scope, input, cond, body, options);
+  }
+
   /**
    * output = cond ? then_branch(input) : else_branch(input)
    *
@@ -6098,10 +6197,29 @@ public StatefulPartitionedCall statefulPartitionedCall(Iterable<Operand<?>> args
    */
   public StatelessIf statelessIf(Operand<? extends TType> cond, Iterable<Operand<?>> input,
       List<Class<? extends TType>> Tout, ConcreteFunction thenBranch, ConcreteFunction elseBranch,
-      StatelessIf.Options... options) {
+      If.Options... options) {
     return StatelessIf.create(scope, cond, input, Tout, thenBranch, elseBranch, options);
   }
 
+  /**
+   * returns {@code f(inputs)}, where {@code f}'s body is placed and partitioned.
+   *
+   * @param args A list of input tensors.
+   * @param Tout A list of output types.
+   * @param f <pre>
+   *    A function that takes 'args', a list of tensors, and returns 'output',
+   *    another list of tensors. Input and output types are specified by 'Tin'
+   *    and 'Tout'. The function body of f will be placed and partitioned across
+   *    devices, setting this op apart from the regular Call op.
+   *  </pre>
+   * @param options carries optional attribute values
+   * @return a new instance of StatelessPartitionedCall
+   */
+  public StatelessPartitionedCall statelessPartitionedCall(Iterable<Operand<?>> args,
+      List<Class<? extends TType>> Tout, ConcreteFunction f, PartitionedCall.Options... options) {
+    return StatelessPartitionedCall.create(scope, args, Tout, f, options);
+  }
+
   /**
    * output = input; While (Cond(output)) { output = Body(output) }
    *
@@ -6127,7 +6245,7 @@ public StatelessIf statelessIf(Operand<? extends TType> cond, Iterable<Operand<?
    * @return a new instance of StatelessWhile
    */
   public StatelessWhile statelessWhile(Iterable<Operand<?>> input, ConcreteFunction cond,
-      ConcreteFunction body, StatelessWhile.Options... options) {
+      ConcreteFunction body, While.Options... options) {
     return StatelessWhile.create(scope, input, cond, body, options);
   }
 
@@ -7990,6 +8108,7 @@ public Where where(Operand<? extends TType> condition) {
 
   /**
    * output = input; While (Cond(output)) { output = Body(output) }
+   *  Selects between {@link StatefulWhile} and {@link StatelessWhile} based on the statefulness of the function arguments.
    *
    * @param input A list of input tensors whose types are T.
    * @param cond <pre>

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core/BatchFunction.java

@@ -109,7 +109,7 @@ public static BatchFunction create(Scope scope, Iterable<Operand<?>> inTensors,
       Iterable<Operand<?>> capturedTensors, ConcreteFunction f, Long numBatchThreads,
       Long maxBatchSize, Long batchTimeoutMicros, List<Class<? extends TType>> Tout,
       Options... options) {
-    OperationBuilder opBuilder = scope.env().opBuilder("BatchFunction", scope.makeOpName("BatchFunction"));
+    OperationBuilder opBuilder = scope.env().opBuilder(OP_NAME, scope.makeOpName("BatchFunction"));
     opBuilder.addInputList(Operands.asOutputs(inTensors));
     opBuilder.addInputList(Operands.asOutputs(capturedTensors));
     opBuilder = scope.apply(opBuilder);

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core/Case.java

@@ -0,0 +1,150 @@
+/* Copyright 2018 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
+
+    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.
+=======================================================================*/
+
+// This class has been generated, DO NOT EDIT!
+
+package org.tensorflow.op.core;
+
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import org.tensorflow.ConcreteFunction;
+import org.tensorflow.Operand;
+import org.tensorflow.Output;
+import org.tensorflow.ndarray.Shape;
+import org.tensorflow.op.Scope;
+import org.tensorflow.op.annotation.Endpoint;
+import org.tensorflow.op.annotation.Operator;
+import org.tensorflow.types.TInt32;
+import org.tensorflow.types.family.TType;
+
+/**
+ * An n-way switch statement which calls a single branch function.
+ * <pre>
+ * An n-way switch statement, implementing the following:
+ * ```
+ * switch (branch_index) {
+ *   case 0:
+ *     output = branches[0](input);
+ *     break;
+ *   case 1:
+ *     output = branches[1](input);
+ *     break;
+ *   ...
+ *   case [[nbranches-1]]:
+ *   default:
+ *     output = branches[nbranches-1](input);
+ *     break;
+ * }
+ * ```
+ * </pre>
+ * Selects between {@link StatefulCase} and {@link StatelessCase} based on the statefulness of the function arguments.
+ */
+@Operator
+public interface Case extends Iterable<Operand<TType>> {
+  /**
+   * Factory method to create a class wrapping a new Case operation.
+   *
+   * @param scope current scope
+   * @param branchIndex The branch selector, an int32 Tensor.
+   * @param input A list of input tensors passed to the branch function.
+   * @param Tout A list of output types.
+   * @param branches <pre>
+   *   A list of functions each of which takes 'inputs' and returns a list of
+   *   tensors, whose types are the same as what every other branch returns.
+   * </pre>
+   * @param options carries optional attribute values
+   * @return a new instance of Case
+   */
+  @Endpoint(
+      describeByClass = true,
+      name = "caseOp"
+  )
+  static Case create(Scope scope, Operand<TInt32> branchIndex, Iterable<Operand<?>> input,
+      List<Class<? extends TType>> Tout, List<ConcreteFunction> branches, Options... options) {
+    boolean isStateful = false;
+    if (branches.stream().anyMatch(x -> x.isStateful())) {
+      isStateful = true;
+    }
+    if (isStateful) {
+      return StatefulCase.create(scope, branchIndex, input, Tout, branches, options);
+    } else {
+      return StatelessCase.create(scope, branchIndex, input, Tout, branches, options);
+    }
+  }
+
+  /**
+   * Sets the outputShapes option.
+   *
+   * @param outputShapes the outputShapes option
+   * @return this Options instance.
+   */
+  static Options outputShapes(List<Shape> outputShapes) {
+    return new Options().outputShapes(outputShapes);
+  }
+
+  /**
+   * Sets the outputShapes option.
+   *
+   * @param outputShapes the outputShapes option
+   * @return this Options instance.
+   */
+  static Options outputShapes(Shape[] outputShapes) {
+    return new Options().outputShapes(outputShapes);
+  }
+
+  /**
+   * Gets output.
+   * A list of return values.
+   * @return output.
+   */
+  List<Output<?>> output();
+
+  @Override
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  Iterator<Operand<TType>> iterator();
+
+  /**
+   * Optional attributes for {@link org.tensorflow.op.core.Case}
+   */
+  class Options {
+    List<Shape> outputShapes;
+
+    private Options() {
+    }
+
+    /**
+     * Sets the outputShapes option.
+     *
+     * @param outputShapes the outputShapes option
+     * @return this Options instance.
+     */
+    public Options outputShapes(List<Shape> outputShapes) {
+      this.outputShapes = outputShapes;
+      return this;
+    }
+
+    /**
+     * Sets the outputShapes option.
+     *
+     * @param outputShapes the outputShapes option
+     * @return this Options instance.
+     */
+    public Options outputShapes(Shape... outputShapes) {
+      this.outputShapes = Arrays.asList(outputShapes);
+      return this;
+    }
+  }
+}

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core/For.java

@@ -78,7 +78,7 @@ private For(Operation operation) {
   )
   public static For create(Scope scope, Operand<TInt32> start, Operand<TInt32> limit,
       Operand<TInt32> delta, Iterable<Operand<?>> input, ConcreteFunction body) {
-    OperationBuilder opBuilder = scope.env().opBuilder("For", scope.makeOpName("For"));
+    OperationBuilder opBuilder = scope.env().opBuilder(OP_NAME, scope.makeOpName("For"));
     opBuilder.addInput(start.asOutput());
     opBuilder.addInput(limit.asOutput());
     opBuilder.addInput(delta.asOutput());

tensorflow-core/tensorflow-core-api/src/gen/java/org/tensorflow/op/core/GroupByReducerDataset.java

@@ -82,7 +82,7 @@ public static GroupByReducerDataset create(Scope scope, Operand<? extends TType>
       Iterable<Operand<?>> finalizeFuncOtherArguments, ConcreteFunction keyFunc,
       ConcreteFunction initFunc, ConcreteFunction reduceFunc, ConcreteFunction finalizeFunc,
       List<Class<? extends TType>> outputTypes, List<Shape> outputShapes) {
-    OperationBuilder opBuilder = scope.env().opBuilder("GroupByReducerDataset", scope.makeOpName("GroupByReducerDataset"));
+    OperationBuilder opBuilder = scope.env().opBuilder(OP_NAME, scope.makeOpName("GroupByReducerDataset"));
     opBuilder.addInput(inputDataset.asOutput());
     opBuilder.addInputList(Operands.asOutputs(keyFuncOtherArguments));
     opBuilder.addInputList(Operands.asOutputs(initFuncOtherArguments));