in tensorflow-core/tensorflow-core-generator/src/main/java/org/tensorflow/processor/operator/OperatorProcessor.java [493:694]
private static TypeSpec buildTopClass(OpsSpec spec) {
// System.out.println("Generating " + spec.className + " class");
MethodSpec.Builder ctorBuilder =
MethodSpec.constructorBuilder()
.addParameter(Names.Scope, "scope")
.addStatement("this.scope = scope", Names.Scope);
TypeSpec.Builder opsBuilder =
TypeSpec.classBuilder("Ops")
.addModifiers(Modifier.PUBLIC, Modifier.FINAL)
.addJavadoc(
"An API for building operations as {@link $T Op}s\n<p>\n"
+ "Any operation wrapper found in the classpath properly annotated as an"
+ "{@link $T @Operator} is exposed\n"
+ "by this API or one of its subgroup.\n<p>Example usage:\n<pre>{@code\n"
+ "try (Graph g = new Graph()) {\n"
+ " Ops tf = Ops.create(g);\n"
+ " // Operations are typed classes with convenience\n"
+ " // builders in Ops.\n"
+ " Constant<TInt32> three = tf.constant(3);\n"
+ " // Single-result operations implement the Operand\n"
+ " // interface, so this works too.\n"
+ " Operand<TInt32> four = tf.constant(4);\n"
+ " // Most builders are found within a group, and accept\n"
+ " // Operand types as operands\n"
+ " Operand<TInt32> nine = tf.math.add(four, tf.constant(5));\n"
+ " // Multi-result operations however offer methods to\n"
+ " // select a particular result for use.\n"
+ " Operand<TInt32> result = \n"
+ " tf.math.add(tf.unique(s, a).y(), b);\n"
+ " // Optional attributes\n"
+ " tf.linalg.matMul(a, b, MatMul.transposeA(true));\n"
+ " // Naming operators\n"
+ " tf.withName(\"foo\").constant(5); // name \"foo\"\n"
+ " // Names can exist in a hierarchy\n"
+ " Ops sub = tf.withSubScope(\"sub\");\n"
+ " sub.withName(\"bar\").constant(4); // \"sub/bar\"\n"
+ "}\n"
+ "}</pre>\n",
Names.Op,
Names.Operator)
.addMethods(spec.methods);
addGroupFields(opsBuilder, ctorBuilder, spec.subGroups, true);
opsBuilder.addMethod(ctorBuilder.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withSubScope")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.String, "childScopeName")
.returns(Names.Ops)
.addStatement("return new $T(scope.withSubScope(childScopeName))", Names.Ops)
.addJavadoc(
"Returns an API that builds operations with the provided name prefix.\n"
+ "\n@see {@link $T#withSubScope(String)}\n",
Names.Scope)
.build());
String initScopeComment =
"<p>\nInit operations will be initialized at session creation, will have their inputs (and control inputs) made init ops as well,\n"
+ "and are ignored when used as control dependencies.\n"
+ "Additionally, this scope ignores any control dependencies.\n"
+ "<p>\nIf an input can not be made an init op (i.e. a Placeholder), will throw an {@link IllegalStateException} on op creation.";
opsBuilder.addMethod(
MethodSpec.methodBuilder("withInitScope")
.addModifiers(Modifier.PUBLIC)
.returns(Names.Ops)
.addStatement("return new $T(scope.withInitScope())", Names.Ops)
.addJavadoc(
"Returns an API that builds init operations. {@link #liftToInitScope(Operand)} will be called for all created operations.\n"
+ initScopeComment
+ "\n@see #liftToInitScope(Operand)")
.build());
TypeVariableName T = TypeVariableName.get("T").withBounds(Names.Operand);
opsBuilder.addMethod(
MethodSpec.methodBuilder("liftToInitScope")
.addTypeVariable(T)
.addModifiers(Modifier.PUBLIC)
.addParameter(T, "op")
.returns(T)
.addStatement("scope.env().registerInitOp(op.op())")
.addStatement("return op")
.addJavadoc(
"Make {@code op} an init operation, doing the same for all of it's inputs (and control inputs).\n"
+ initScopeComment
+ "\n@see ExecutionEnvironment#registerInitOp(Operation)\n"
+ "\n@throws IllegalStateException if the op or one of its inputs can't be made an init op.")
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withName")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.String, "opName")
.returns(Names.Ops)
.addStatement("return new Ops(scope.withName(opName))")
.addJavadoc(
"Returns an API that uses the provided name for an op.\n\n"
+ "@see {@link $T#withName(String)}\n",
Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withDevice")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.DeviceSpec, "deviceSpec")
.returns(Names.Ops)
.addStatement("return new Ops(scope.withDevice(deviceSpec))")
.addJavadoc(
"Returns an API that places the created operations on the device(s) matching the provided spec.\n\n"
+ "@see {@link $T#withDevice(DeviceSpec)}\n",
Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withControlDependencies")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.IterableOp, "controls")
.returns(Names.Ops)
.addStatement("return new Ops(scope.withControlDependencies(controls))")
.addJavadoc(
"Returns an API that adds operations to the graph with the provided control dependencies.\n\n"
+ "@see {@link $T#withControlDependencies(Iterable<Op<?>>)}\n",
Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withControlDependencies")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.ArrayOp, "controls")
.varargs()
.returns(Names.Ops)
.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",
Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withControlDependencyOps")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.IterableOperation, "controls")
.returns(Names.Ops)
.addStatement("return new Ops(scope.withControlDependencyOps(controls))")
.addJavadoc(
"Returns an API that adds operations to the graph with the provided control dependencies.\n\n"
+ "@see {@link $T#withControlDependencyOps(Iterable<Operation>)}\n",
Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("withControlDependencyOps")
.addModifiers(Modifier.PUBLIC)
.addParameter(Names.ArrayOperation, "controls")
.varargs()
.returns(Names.Ops)
.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",
Names.Scope)
.build());
opsBuilder.addField(
FieldSpec.builder(Names.Scope, "scope")
.addModifiers(Modifier.PRIVATE, Modifier.FINAL)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("scope")
.addModifiers(Modifier.PUBLIC, Modifier.FINAL)
.returns(Names.Scope)
.addStatement("return scope")
.addJavadoc("Returns the current {@link $T scope} of this API\n", Names.Scope)
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("create")
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
.addParameter(Names.ExecutionEnvironment, "env")
.returns(Names.Ops)
.addStatement("return new Ops(env.baseScope())")
.addJavadoc(
"Creates an API for building operations in the provided execution environment\n")
.build());
opsBuilder.addMethod(
MethodSpec.methodBuilder("create")
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
.returns(Names.Ops)
.addStatement("return create($T.getDefault())", Names.EagerSession)
.addJavadoc(
"Creates an API for building operations in the default eager execution environment\n\n"
+ "<p>Invoking this method is equivalent to {@code Ops.create(EagerSession.getDefault())}.\n")
.build());
return opsBuilder.build();
}