Proper attribute setters

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

Author: rnett

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

@@ -238,7 +238,7 @@ public Map<String, Operand<?>> call(Scope scope,
 
     opBuilder.addInputList(inputList.stream().map(Operand::asOutput).toArray(Output[]::new));
 
-    opBuilder.setFunctionName("f", name);
+    opBuilder.setAttr("f", this);
     opBuilder.setAttr("Tin", inputList.stream().map(x -> x.asOutput().dataType()).toArray(DataType[]::new));
     opBuilder.setAttr("Tout", signature().getOutputs().values().stream().map(x -> x.dataType).toArray(DataType[]::new));
 

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

@@ -22,6 +22,7 @@
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrBoolList;
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrFloat;
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrFloatList;
+import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrFunctionList;
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrFunctionName;
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrInt;
 import static org.tensorflow.internal.c_api.global.tensorflow.TFE_OpSetAttrIntList;
@@ -36,6 +37,9 @@
 
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.bytedeco.javacpp.BooleanPointer;
 import org.bytedeco.javacpp.BytePointer;
 import org.bytedeco.javacpp.IntPointer;
@@ -219,8 +223,22 @@ public EagerOperationBuilder setAttr(String name, Shape[] values) {
   }
 
   @Override
-  public OperationBuilder setFunctionName(String attrName, String functionName) {
-    setAttrFunctionName(opHandle, attrName, functionName);
+  public OperationBuilder setAttr(String name, ConcreteFunction value) {
+    session.attachFunction(value);
+    setAttrFunctionName(opHandle, name, value.getNativeFunctionName());
+    return this;
+  }
+
+  @Override
+  public OperationBuilder setAttr(String name, ConcreteFunction[] value) {
+    for (ConcreteFunction fn : value) {
+      session.attachFunction(fn);
+    }
+
+    setAttrFunctionList(opHandle, session.nativeHandle(), name, Arrays.stream(value)
+        .map(ConcreteFunction::getNativeFunctionName)
+        .collect(Collectors.toList()));
+
     return this;
   }
 
@@ -416,7 +434,7 @@ private static void setAttrShapeList(TFE_Op opHandle, String name, long[] shapes
       }
       TF_Status status = TF_Status.newStatus();
       TFE_OpSetAttrShapeList(opHandle, new BytePointer(name), shapesPointers, new IntPointer(numDims),
-              numDims.length, status);
+          numDims.length, status);
     }
   }
 
@@ -426,4 +444,20 @@ private static void setAttrFunctionName(TFE_Op opHandle, String attrName, String
       TFE_OpSetAttrFunctionName(opHandle, attrName, functionName, functionName.length());
     }
   }
+
+  private static void setAttrFunctionList(TFE_Op opHandle, TFE_Context context, String attrName,
+      List<String> functionNames) {
+    requireOp(opHandle);
+    requireContext(context);
+    try (PointerScope scope = new PointerScope()) {
+      PointerPointer<TFE_Op> fns = new PointerPointer<>(functionNames.size());
+      for (int i = 0; i < functionNames.size(); i++) {
+        TF_Status status = TF_Status.newStatus();
+        TFE_Op op = TFE_Op.newOp(context, functionNames.get(i), status);
+        status.throwExceptionIfNotOK();
+        fns.put(i, op);
+      }
+      TFE_OpSetAttrFunctionList(opHandle, new BytePointer(attrName), fns, functionNames.size());
+    }
+  }
 }

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

@@ -35,9 +35,13 @@
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrTensorList;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrType;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrTypeList;
+import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetAttrValueProto;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetDevice;
 
 import java.nio.charset.Charset;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.bytedeco.javacpp.BooleanPointer;
 import org.bytedeco.javacpp.BytePointer;
 import org.bytedeco.javacpp.IntPointer;
@@ -54,7 +58,10 @@
 import org.tensorflow.internal.c_api.TF_Status;
 import org.tensorflow.internal.c_api.TF_Tensor;
 import org.tensorflow.ndarray.Shape;
+import org.tensorflow.proto.framework.AttrValue;
+import org.tensorflow.proto.framework.AttrValue.ListValue;
 import org.tensorflow.proto.framework.DataType;
+import org.tensorflow.proto.framework.NameAttrList;
 
 /**
  * An {@link OperationBuilder} for adding {@link GraphOperation}s to a {@link Graph}.
@@ -347,9 +354,24 @@ public GraphOperationBuilder setAttr(String name, String[] value) {
   }
 
   @Override
-  public OperationBuilder setFunctionName(String attrName, String functionName) {
+  public OperationBuilder setAttr(String name, ConcreteFunction value) {
+    graph.attachFunction(value);
     try (Reference r = graph.ref()) {
-      setAttrFunctionName(unsafeNativeHandle, attrName, functionName);
+      setAttrFunctionName(unsafeNativeHandle, name, value.getNativeFunctionName());
+    }
+    return this;
+  }
+
+  @Override
+  public OperationBuilder setAttr(String name, ConcreteFunction[] value) {
+    for (ConcreteFunction f : value) {
+      graph.attachFunction(f);
+    }
+
+    try (Reference r = graph.ref()) {
+      setAttrFunctionList(unsafeNativeHandle, name, Arrays.stream(value)
+          .map(ConcreteFunction::getNativeFunctionName)
+          .collect(Collectors.toList()));
     }
     return this;
   }
@@ -556,4 +578,20 @@ private static void setAttrFunctionName(TF_OperationDescription opHandle, String
       TF_SetAttrFuncName(opHandle, attrName, functionName, functionName.length());
     }
   }
+
+  private static void setAttrFunctionList(TF_OperationDescription opHandle, String attrName,
+      List<String> functionNames) {
+    requireHandle(opHandle);
+    try (PointerScope scope = new PointerScope()) {
+      TF_Status status = TF_Status.newStatus();
+      AttrValue value = AttrValue.newBuilder().setList(ListValue.newBuilder().addAllFunc(
+          functionNames.stream()
+              .map(x -> NameAttrList.newBuilder().setName(x).build())
+              .collect(Collectors.toList())
+      ).build()).build();
+      byte[] bytes = value.toByteArray();
+      TF_SetAttrValueProto(opHandle, attrName, new BytePointer(bytes), bytes.length, status);
+      status.throwExceptionIfNotOK();
+    }
+  }
 }

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

@@ -226,11 +226,22 @@ public interface OperationBuilder {
   OperationBuilder setAttr(String name, Shape[] value);
 
   /**
-   * Set a function name attribute of the operation being build.
+   * Set the function value of an attribute of the operation being built. Also attaches the function and dependencies to
+   * the execution environment.
    *
-   * @param attrName the attribute to set
-   * @param functionName the function name
+   * @param name attribute name
+   * @param value attribute value
+   * @return the OperationBuilder instance for chaining.
+   */
+  OperationBuilder setAttr(String name, ConcreteFunction value);
+
+  /**
+   * Set the function values of an attribute of the operation being built. Also attaches the functions and dependencies
+   * to the execution environment.
+   *
+   * @param name attribute name
+   * @param value attribute value
    * @return the OperationBuilder instance for chaining.
    */
-  OperationBuilder setFunctionName(String attrName, String functionName);
+  OperationBuilder setAttr(String name, ConcreteFunction[] value);
 }

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

@@ -124,7 +124,7 @@ public void setAttrs() {
           .build();
       // bool
       opBuilder(session, "All", "Bool")
-          .addInput(tf.constant(new boolean[] {true, true, false}).asOutput())
+          .addInput(tf.constant(new boolean[]{true, true, false}).asOutput())
           .addInput(tf.constant(0).asOutput())
           .setAttr("keep_dims", false)
           .build();
@@ -134,7 +134,8 @@ public void setAttrs() {
           .addInput(tf.constant(10.00000f).asOutput())
           .setAttr("tolerance", 0.1f)
           .build();
-      // Missing tests: list(string), list(byte), list(bool), list(type)
+      // Missing tests: list(string), list(byte), list(bool), list(type), list(func)
+      // func is done via ConcreteFunction execution
     }
   }
 

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

@@ -96,17 +96,18 @@ public void setAttr() {
       g.opBuilder("MaxPool", "IntList")
           .addInput(tf.constant(new float[2][2][2][2]).asOutput())
           .setAttr("ksize", new long[] {1, 1, 1, 1})
-          .setAttr("strides", new long[] {1, 1, 1, 1})
+          .setAttr("strides", new long[]{1, 1, 1, 1})
           .setAttr("padding", "SAME")
           .build();
       assertTrue(hasNode(g, "IntList"));
       // list(float)
       g.opBuilder("FractionalMaxPool", "FloatList")
           .addInput(tf.constant(new float[2][2][2][2]).asOutput())
-          .setAttr("pooling_ratio", new float[] {1.0f, 1.44f, 1.73f, 1.0f})
+          .setAttr("pooling_ratio", new float[]{1.0f, 1.44f, 1.73f, 1.0f})
           .build();
       assertTrue(hasNode(g, "FloatList"));
-      // Missing tests: float, list(dtype), list(tensor), list(string), list(bool)
+      // Missing tests: float, list(dtype), list(tensor), list(string), list(bool), list(func)
+      // func is done via ConcreteFunction execution
     }
   }