chore: import upstream snapshot with attribution

tensorflow/go/genop/.gitignore

@@ -0,0 +1,2 @@
+# .pb.go files generated by generate.sh
+internal/proto/*

tensorflow/go/genop/BUILD

@@ -0,0 +1 @@
+# Empty file to be replaced in https://github.com/tensorflow/tensorflow/pull/50934

tensorflow/go/genop/generate.go

@@ -0,0 +1,21 @@
+/*
+Copyright 2020 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.
+*/
+
+// Package stub file for non-Windows builds.
+
+//go:generate bash generate.sh
+
+package main

tensorflow/go/genop/generate.sh

@@ -0,0 +1,77 @@
+#!/usr/bin/env bash
+# Copyright 2016 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.
+# ==============================================================================
+
+set -e
+
+go get -d google.golang.org/protobuf/proto
+go install google.golang.org/protobuf/cmd/protoc-gen-go@latest
+
+if [ -z "${GOPATH}" ]
+then
+  GOPATH=$(go env GOPATH)
+fi
+
+if [ -z "${GOOS}" ]
+then
+  GOOS=$(go env GOOS)
+fi
+
+# convert GOPATH's Windows style to UNIX style
+if [[ $GOOS == "windows" ]]; then
+  # eg: convert "D:\go-14;D:\go-13" to "D\go-14;D\go-13"
+  GOPATH=${GOPATH//:\\/\\}
+  # eg: convert "D\go-14;D\go-13" to "\D\go-14:\D\go-13"
+  GOPATH=\\${GOPATH//;/:\\}
+  # eg: convert "\D\go-14:\D\go-13" to "/D/go-14:/D/go-13"
+  GOPATH=${GOPATH//\\/\/}
+fi
+
+cd $(dirname $0)
+for g in $(echo "${GOPATH//:/ }"); do
+    TF_DIR="${g}/src/github.com/tensorflow/tensorflow"
+    PROTOC="${TF_DIR}/bazel-out/host/bin/external/protobuf/protoc"
+    if [ -x "${PROTOC}" ]; then
+        break
+    fi
+done
+
+if [ ! -x "${PROTOC}" ]
+then
+  set +e
+  PATH_PROTOC=$(which protoc)
+  if [ ! -x "${PATH_PROTOC}" ]
+  then
+    echo "Protocol buffer compiler protoc not found in PATH or in ${PROTOC}"
+    echo "Perhaps build it using:"
+    echo "bazel build --config opt @com_google_protobuf//:protoc"
+    exit 1
+  fi
+  PROTOC=$PATH_PROTOC
+  set -e
+fi
+
+# Ensure that protoc-gen-go is available in $PATH
+# Since ${PROTOC} will require it.
+export PATH=$PATH:${GOPATH}/bin
+for FILE in ${TF_DIR}/tensorflow/core/framework/*.proto \
+    ${TF_DIR}/tensorflow/core/protobuf/*.proto \
+    ${TF_DIR}/tensorflow/compiler/xla/pjrt/distributed/*.proto \
+    ${TF_DIR}/tensorflow/compiler/xla/stream_executor/*.proto; do
+  ${PROTOC} \
+    -I ${TF_DIR} \
+    --go_out=${GOPATH}/src \
+    $FILE
+done

tensorflow/go/genop/internal/BUILD

@@ -0,0 +1 @@
+# Empty file to be replaced in https://github.com/tensorflow/tensorflow/pull/50934

tensorflow/go/genop/internal/api_def_map.go

@@ -0,0 +1,129 @@
+/*
+Copyright 2017 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.
+*/
+
+package internal
+
+/*
+#include <stdlib.h>
+#include <string.h>
+
+#include "tensorflow/c/c_api.h"
+*/
+import "C"
+
+import (
+	"errors"
+	"fmt"
+	"runtime"
+	"unsafe"
+
+	"google.golang.org/protobuf/proto"
+
+	adpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/api_def_go_proto"
+	odpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/op_def_go_proto"
+)
+
+// Encapsulates a collection of API definitions.
+//
+// apiDefMap represents a map from operation name to corresponding
+// ApiDef proto (see
+// https://www.tensorflow.org/code/tensorflow/core/framework/api_def.proto
+// for ApiDef proto definition).
+type apiDefMap struct {
+	c *C.TF_ApiDefMap
+}
+
+// Creates and returns a new apiDefMap instance.
+//
+// oplist is and OpList proto instance (see
+// https://www.tensorflow.org/code/tensorflow/core/framework/op_def.proto
+// for OpList proto definition).
+
+func newAPIDefMap(oplist *odpb.OpList) (*apiDefMap, error) {
+	// Create a buffer containing the serialized OpList.
+	opdefSerialized, err := proto.Marshal(oplist)
+	if err != nil {
+		return nil, fmt.Errorf("could not serialize OpDef for %s", oplist.String())
+	}
+	data := C.CBytes(opdefSerialized)
+	defer C.free(data)
+
+	opbuf := C.TF_NewBuffer()
+	defer C.TF_DeleteBuffer(opbuf)
+	opbuf.data = data
+	opbuf.length = C.size_t(len(opdefSerialized))
+
+	// Create ApiDefMap.
+	status := C.TF_NewStatus()
+	defer C.TF_DeleteStatus(status)
+	capimap := C.TF_NewApiDefMap(opbuf, status)
+	if C.TF_GetCode(status) != C.TF_OK {
+		return nil, errors.New(C.GoString(C.TF_Message(status)))
+	}
+	apimap := &apiDefMap{capimap}
+	runtime.SetFinalizer(
+		apimap,
+		func(a *apiDefMap) {
+			C.TF_DeleteApiDefMap(a.c)
+		})
+	return apimap, nil
+}
+
+// Updates apiDefMap with the overrides specified in `data`.
+//
+// data - ApiDef text proto.
+func (m *apiDefMap) Put(data string) error {
+	cdata := C.CString(data)
+	defer C.free(unsafe.Pointer(cdata))
+	status := C.TF_NewStatus()
+	defer C.TF_DeleteStatus(status)
+	C.TF_ApiDefMapPut(m.c, cdata, C.size_t(len(data)), status)
+	if C.TF_GetCode(status) != C.TF_OK {
+		return errors.New(C.GoString(C.TF_Message(status)))
+	}
+	return nil
+}
+
+// Returns ApiDef proto instance for the TensorFlow operation
+// named `opname`.
+func (m *apiDefMap) Get(opname string) (*adpb.ApiDef, error) {
+	cname := C.CString(opname)
+	defer C.free(unsafe.Pointer(cname))
+	status := C.TF_NewStatus()
+	defer C.TF_DeleteStatus(status)
+	apidefBuf := C.TF_ApiDefMapGet(
+		m.c, cname, C.size_t(len(opname)), status)
+	defer C.TF_DeleteBuffer(apidefBuf)
+	if C.TF_GetCode(status) != C.TF_OK {
+		return nil, errors.New(C.GoString(C.TF_Message(status)))
+	}
+	if apidefBuf == nil {
+		return nil, fmt.Errorf("could not find ApiDef for %s", opname)
+	}
+
+	var (
+		apidef = new(adpb.ApiDef)
+		size   = int(apidefBuf.length)
+		// A []byte backed by C memory.
+		// See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
+		data = (*[1 << 30]byte)(unsafe.Pointer(apidefBuf.data))[:size:size]
+		err  = proto.Unmarshal(data, apidef)
+	)
+	if err != nil {
+		return nil, err
+	}
+	return apidef, nil
+}

tensorflow/go/genop/internal/genop.go

@@ -0,0 +1,609 @@
+/*
+Copyright 2016 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.
+*/
+
+// Package internal generates Go source code with functions for TensorFlow operations.
+//
+// The basic outline of the generated API is as follows:
+//
+// - One function for each TensorFlow operation
+// - The arguments to the function are the inputs and required attributes of the operation
+// - The function returns the outputs
+// - A function is also generated for each optional attribute of the operation.
+//
+// There is a possibility that there are name collisions between the functions
+// generated for ops and the functions generated for optional attributes. For
+// now, we ignore those, but will need to revisit if a collision is actually
+// encountered.
+package internal
+
+/*
+#include <stdlib.h>
+
+#include "tensorflow/c/c_api.h"
+*/
+import "C"
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path"
+	"reflect"
+	"sort"
+	"strings"
+	"text/template"
+	"unsafe"
+
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
+
+	adpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/api_def_go_proto"
+	odpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/op_def_go_proto"
+)
+
+// GenerateFunctionsForRegisteredOps writes a Go source code file to w
+// containing functions for each TensorFlow operation registered in the address
+// space of the calling process.
+// apidefDirs should be a contain of directories containing api_def_*.pbtxt
+// files to load.
+func GenerateFunctionsForRegisteredOps(
+	w io.Writer, apidefDirs []string) error {
+	ops, apimap, err := registeredOps()
+	if err != nil {
+		return err
+	}
+	for _, dir := range apidefDirs {
+		if err = updateAPIDefs(apimap, dir); err != nil {
+			return err
+		}
+	}
+	return generateFunctionsForOps(w, ops, apimap)
+}
+
+func registeredOps() (*odpb.OpList, *apiDefMap, error) {
+	buf := C.TF_GetAllOpList()
+	defer C.TF_DeleteBuffer(buf)
+	var (
+		list = new(odpb.OpList)
+		size = int(buf.length)
+		// A []byte backed by C memory.
+		// See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
+		data = (*[1 << 30]byte)(unsafe.Pointer(buf.data))[:size:size]
+		err  = proto.Unmarshal(data, list)
+	)
+	if err != nil {
+		return nil, nil, err
+	}
+	// Sort ops by name
+	sort.Slice(list.Op, func(i, j int) bool {
+		return list.Op[i].Name < list.Op[j].Name
+	})
+	apimap, err := newAPIDefMap(list)
+	return list, apimap, err
+}
+
+func updateAPIDefs(m *apiDefMap, dir string) error {
+	files, err := os.ReadDir(dir)
+	if err != nil {
+		return err
+	}
+	for _, file := range files {
+		if file.IsDir() || !strings.HasSuffix(file.Name(), ".pbtxt") {
+			continue
+		}
+		data, err := os.ReadFile(path.Join(dir, file.Name()))
+		if err != nil {
+			return fmt.Errorf("failed to read %q: %v", file.Name(), err)
+		}
+		if err = m.Put(string(data)); err != nil {
+			return fmt.Errorf("failed to process %q: %v", file.Name(), err)
+		}
+	}
+	return nil
+}
+
+func generateFunctionsForOps(w io.Writer, ops *odpb.OpList, apimap *apiDefMap) error {
+	thisPackage := reflect.TypeOf(tmplArgs{}).PkgPath()
+	if err := tmplHeader.Execute(w, thisPackage); err != nil {
+		return err
+	}
+	denylist := map[string]bool{
+		"Const":           true,
+		"PyFunc":          true,
+		"PyFuncStateless": true,
+	}
+	for _, op := range ops.Op {
+		if denylist[op.Name] {
+			continue
+		}
+		apidef, err := apimap.Get(op.Name)
+		if err != nil {
+			return err
+		}
+		if err := generateFunctionForOp(w, op, apidef); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func generateFunctionForOp(w io.Writer, op *odpb.OpDef, apidef *adpb.ApiDef) error {
+	if strings.HasPrefix(op.Name, "_") { // Internal operation
+		return nil
+	}
+	// Ignore operations where the Go types corresponding to the TensorFlow
+	// type haven't been worked out (such as "func"s).
+	for _, a := range op.Attr {
+		if _, err := goType(a.Type); err != nil {
+			return nil
+		}
+	}
+	// Also, haven't figured out reference types yet, so ignore those too.
+	for _, a := range op.InputArg {
+		if a.IsRef {
+			return nil
+		}
+	}
+	for _, a := range op.OutputArg {
+		if a.IsRef {
+			return nil
+		}
+	}
+	if apidef.Summary == "" {
+		// Undocumented operation, perhaps a sign of not being ready to
+		// export.
+		return nil
+	}
+	tmplArgs, err := newTmplArgs(op, apidef)
+	if err != nil {
+		return err
+	}
+	return tmplOp.Execute(w, tmplArgs)
+}
+
+var (
+	// Go keywords that cannot be used as identifiers.
+	// From https://golang.org/ref/spec#Keywords
+	keywords = []string{
+		"break", "default", "func", "interface", "select", "case",
+		"defer", "go", "map", "struct", "chan", "else", "goto",
+		"package", "switch", "const", "fallthrough", "if", "range",
+		"type", "continue", "for", "import", "return", "var",
+	}
+
+	tmplHeader = template.Must(template.New("header").Parse(`// DO NOT EDIT
+// This file was machine generated by {{.}}
+//
+// WARNING: This generation of wrapper function for TensorFlow ops is in an
+// experimental state. The generated API can change without notice.
+
+package op
+
+import tf "github.com/tensorflow/tensorflow/tensorflow/go"
+
+// optionalAttr is an intentionally un-exported type to hide
+// details of how optional attributes to operations are implemented.
+type optionalAttr map[string]interface{}
+
+func makeOutputList(op *tf.Operation, start int, output string) ([]tf.Output, int, error) {
+	size, err := op.OutputListSize(output)
+	if err != nil {
+		return nil, start, err
+	}
+	list := make([]tf.Output, size)
+	for i := 0; i < size; i++ {
+		list[i] = op.Output(start + i)
+	}
+	return list, start + size, nil
+}
+`))
+
+	tmplOp = template.Must(template.New("op").Funcs(template.FuncMap{
+		"MakeComment":         makeComment,
+		"GoType":              goType,
+		"CamelCase":           camelCase,
+		"Identifier":          identifier,
+		"IsListArg":           isListArg,
+		"IsListAttr":          isListAttr,
+		"MarshalProtoMessage": marshalProtoMessage,
+	}).Parse(`
+{{if .OptionalAttrs -}}
+{{/* Type for specifying all optional attributes. */ -}}
+// {{.Op.Name}}Attr is an optional argument to {{.Op.Name}}.
+type {{.Op.Name}}Attr func(optionalAttr)
+
+{{range .OptionalAttrs}}
+// {{$.Op.Name}}{{CamelCase .RenameTo}} sets the optional {{.RenameTo}} attribute to value.
+{{- if .Description}}
+//
+// value: {{MakeComment .Description}}
+{{- end}}
+// If not specified, defaults to {{MarshalProtoMessage .DefaultValue}}
+{{- if .HasMinimum}}
+//
+// {{if .IsListAttr }}REQUIRES: len(value) >= {{.Minimum}}{{else}}REQUIRES: value >= {{.Minimum}}{{end}}
+{{- end}}
+func {{$.Op.Name}}{{CamelCase .RenameTo}}(value {{GoType .Type}}) {{$.Op.Name}}Attr {
+	return func(m optionalAttr) {
+		m[{{printf "%q" .Name}}] = value
+	}
+}
+{{end}}
+{{end}}
+
+{{- /* Create a godoc friendly comment. */ -}}
+
+// {{MakeComment .APIDef.Summary}}
+
+{{- with .Op.Deprecation}}
+//
+// DEPRECATED at GraphDef version {{.Version}}: {{.Explanation}}
+{{- end -}}
+
+{{- with .APIDef.Description}}
+//
+// {{MakeComment .}}
+{{- end -}}
+
+{{- if .DescribeArguments}}
+//
+// Arguments:
+{{- range .InArgsReordered}}
+//	{{if .Description}}{{Identifier .RenameTo}}: {{MakeComment .Description}}{{end}}
+{{- end -}}
+{{- range .RequiredAttrs}}
+//	{{if .Description}}{{Identifier .RenameTo}}: {{MakeComment .Description}}{{end}}
+{{- end -}}
+{{- end -}}
+
+{{- if (not .Op.OutputArg) }}
+//
+// Returns the created operation.
+{{- else }}
+{{- if .DescribeOutputs}}
+//
+{{- if eq (len .OutArgs) 1 }}
+// Returns {{range .OutArgs}}{{MakeComment .Description}}{{end}}
+{{- else }}
+// Returns:
+{{- range .OutArgs}}
+//	{{Identifier .RenameTo}}{{if .Description}}: {{MakeComment .Description}}{{end}}
+{{- end -}}
+{{- end -}}
+{{- end -}}
+{{- end -}}
+{{- /*
+
+  The function signature.
+  Since OpDef.Name is in CamelCase, it cannot conflict with a reserved keyword in Golang
+*/}}
+func {{.Op.Name}}
+
+{{- /*
+  Fill in input arguments:
+  (1) The Scope
+  (2) All input arguments (which may be either []tf.Output or tf.Output)
+  (3) All required attributes
+  (4) Variadic list of optional attributes
+*/ -}}
+
+(scope *Scope
+{{- range $i, $a := .InArgsReordered}}, {{Identifier $a.RenameTo}} {{if $a.IsListArg}}[]{{end}}tf.Output{{end -}}
+{{range $i, $a := .RequiredAttrs}}, {{Identifier $a.RenameTo}} {{GoType $a.Type}}{{end -}}
+{{if .OptionalAttrs}}, optional ...{{.Op.Name}}Attr{{end -}}
+)
+
+{{- /* Construct outputs: len(.OutArgs) or a *tf.Operation */ -}}
+
+{{if .OutArgs -}}
+({{range $i,$a := .OutArgs}}{{if $i}}, {{end}}{{Identifier $a.RenameTo}} {{if $a.IsListArg}}[]{{end}}tf.Output{{end -}})
+{{- else -}}
+(o *tf.Operation)
+{{- end }} {
+	if scope.Err() != nil {
+		return
+	}
+	{{if .HasAttrs -}}
+	attrs := map[string]interface{}{ {{- range .RequiredAttrs}}{{printf "%q" .Name}}: {{Identifier .RenameTo}},{{end}}}
+	{{if .OptionalAttrs -}}
+	for _, a := range optional {
+		a(attrs)
+	}
+	{{end -}}
+	{{end -}}
+	opspec := tf.OpSpec{
+		Type: {{printf "%q" .Op.Name}},
+		{{if .InArgs -}}
+		Input: []tf.Input{
+			{{range $i,$a := .InArgs}}{{if $a.IsListArg}}tf.OutputList({{Identifier $a.RenameTo}}){{else}}{{Identifier $a.RenameTo}}{{end}}, {{end}}
+		},
+		{{- end}}
+		{{- if .HasAttrs}}
+		Attrs: attrs,
+		{{- end}}
+	}
+	{{- if .OutArgs}}
+	{{- if .HasListOutput}}
+	op := scope.AddOperation(opspec)
+	if scope.Err() != nil {
+		return
+	}
+	var idx int
+	var err error
+	{{- range $i, $a := .OutArgs}}
+	{{- if $a.IsListArg}}
+	if {{Identifier .RenameTo}}, idx, err = makeOutputList(op, idx, {{printf "%q" .Name}}); err != nil {
+		scope.UpdateErr({{printf "%q" $.Op.Name}}, err)
+		return
+	}
+	{{- else }}
+	{{Identifier .RenameTo}} = op.Output(idx)
+	{{- end }}{{- /* if IsListArg */}}
+	{{- end }}{{- /* range .OutArgs */}}
+	return {{range $i, $a := .OutArgs}}{{if $i}}, {{end}}{{Identifier .RenameTo}}{{end}}
+	{{- else }}
+	op := scope.AddOperation(opspec)
+	return {{range $i, $a := .OutArgs}}{{if $i}}, {{end}}op.Output({{$i}}){{end}}
+	{{- end }}{{- /* if .HasListOutput */}}
+	{{- else }}
+	return scope.AddOperation(opspec)
+	{{- end }}{{- /* if .OutArgs */}}
+}
+`))
+)
+
+type attrWrapper struct {
+	op  *odpb.OpDef_AttrDef
+	api *adpb.ApiDef_Attr
+}
+
+func (a *attrWrapper) Name() string              { return a.api.Name }
+func (a *attrWrapper) RenameTo() string          { return a.api.RenameTo }
+func (a *attrWrapper) Description() string       { return a.api.Description }
+func (a *attrWrapper) Type() string              { return a.op.Type }
+func (a *attrWrapper) IsListAttr() bool          { return isListAttr(a.op) }
+func (a *attrWrapper) HasMinimum() bool          { return a.op.HasMinimum }
+func (a *attrWrapper) Minimum() int64            { return a.op.Minimum }
+func (a *attrWrapper) DefaultValue() interface{} { return a.api.DefaultValue }
+
+type argWrapper struct {
+	op  *odpb.OpDef_ArgDef
+	api *adpb.ApiDef_Arg
+}
+
+func (a *argWrapper) Name() string        { return a.api.Name }
+func (a *argWrapper) RenameTo() string    { return a.api.RenameTo }
+func (a *argWrapper) Description() string { return a.api.Description }
+func (a *argWrapper) IsListArg() bool     { return isListArg(a.op) }
+
+type tmplArgs struct {
+	Op     *odpb.OpDef
+	APIDef *adpb.ApiDef
+	// Op.Attr is split into two categories
+	// (1) Required: These must be specified by the client and are thus
+	//     included in the function signature.
+	// (2) Optional: These need not be specified (as they have default
+	//     values) and thus do not appear in the function signature.
+	RequiredAttrs []*attrWrapper
+	OptionalAttrs []*attrWrapper
+	InArgs        []*argWrapper
+	// Input arguments ordered based on arg_order field of ApiDef.
+	InArgsReordered []*argWrapper
+	OutArgs         []*argWrapper
+}
+
+func newTmplArgs(op *odpb.OpDef, apidef *adpb.ApiDef) (*tmplArgs, error) {
+	ret := tmplArgs{Op: op, APIDef: apidef}
+
+	// Setup InArgs field
+	for i, in := range op.InputArg {
+		argCombined := argWrapper{op: in, api: apidef.InArg[i]}
+		ret.InArgs = append(ret.InArgs, &argCombined)
+	}
+
+	// Setup OutArgs field
+	for i, out := range op.OutputArg {
+		argCombined := argWrapper{op: out, api: apidef.OutArg[i]}
+		ret.OutArgs = append(ret.OutArgs, &argCombined)
+	}
+
+	// Setup InArgsReordered field
+	for _, argName := range apidef.ArgOrder {
+		// Find the argument in op.InputArg
+		argIndex := -1
+		for i, in := range op.InputArg {
+			if in.Name == argName {
+				argIndex = i
+				break
+			}
+		}
+		if argIndex == -1 {
+			return nil, fmt.Errorf(
+				"couldn't find argument %s in ApiDef for op %s",
+				argName, op.Name)
+		}
+		argCombined := argWrapper{
+			op: op.InputArg[argIndex], api: apidef.InArg[argIndex]}
+		ret.InArgsReordered = append(ret.InArgsReordered, &argCombined)
+	}
+
+	if len(op.Attr) == 0 {
+		return &ret, nil
+	}
+	// Attributes related to the InputArg's type are inferred automatically
+	// and are not exposed to the client.
+	inferred := make(map[string]bool)
+	for _, in := range op.InputArg {
+		switch {
+		case in.TypeAttr != "":
+			inferred[in.TypeAttr] = true
+		case in.TypeListAttr != "":
+			inferred[in.TypeListAttr] = true
+		}
+		if in.NumberAttr != "" {
+			inferred[in.NumberAttr] = true
+		}
+	}
+	for i, attr := range op.Attr {
+		if inferred[attr.Name] {
+			continue
+		}
+		attrCombined := attrWrapper{op: attr, api: apidef.Attr[i]}
+		if attr.DefaultValue == nil {
+			ret.RequiredAttrs = append(ret.RequiredAttrs, &attrCombined)
+		} else {
+			ret.OptionalAttrs = append(ret.OptionalAttrs, &attrCombined)
+		}
+	}
+	return &ret, nil
+}
+
+func (a *tmplArgs) HasAttrs() bool { return len(a.RequiredAttrs)+len(a.OptionalAttrs) > 0 }
+func (a *tmplArgs) DescribeArguments() bool {
+	for _, arg := range a.InArgs {
+		if arg.Description() != "" {
+			return true
+		}
+	}
+	for _, attr := range a.RequiredAttrs {
+		if attr.Description() != "" {
+			return true
+		}
+	}
+	return false
+
+}
+func (a *tmplArgs) DescribeOutputs() bool {
+	for _, arg := range a.OutArgs {
+		if arg.Description() != "" {
+			return true
+		}
+	}
+	return false
+}
+func (a *tmplArgs) HasListOutput() bool {
+	for _, arg := range a.OutArgs {
+		if arg.IsListArg() {
+			return true
+		}
+	}
+	return false
+}
+
+func makeComment(lines string) string {
+	return strings.Join(strings.SplitAfter(lines, "\n"), "// ")
+}
+
+// goType converts a TensorFlow "type" ('string', 'int', 'list(string)' etc.)
+// to the corresponding type in Go.
+func goType(tfType string) (string, error) {
+	list, tfType := parseTFType(tfType)
+	var gotype string
+	switch tfType {
+	case "int":
+		gotype = "int64"
+	case "float":
+		gotype = "float32"
+	case "bool":
+		gotype = "bool"
+	case "type":
+		gotype = "tf.DataType"
+	case "shape":
+		gotype = "tf.Shape"
+	case "tensor":
+		gotype = "tf.Tensor"
+	case "string":
+		gotype = "string"
+	default:
+		return "", fmt.Errorf("%q is not a recognized DataType", tfType)
+	}
+	if list {
+		gotype = "[]" + gotype
+	}
+	return gotype, nil
+}
+
+func camelCase(snakeCase string) string {
+	words := strings.Split(snakeCase, "_")
+	for i, w := range words {
+		words[i] = strings.ToUpper(string(w[0])) + w[1:]
+	}
+	return strings.Join(words, "")
+}
+
+// identifier creates an identifier for s usable in the generated Go source
+// code.
+//
+// Avoids collisions with keywords and other identifiers used in the generated
+// code.
+func identifier(s string) string {
+	// Identifiers used in the generated code.
+	if s == "tf" || s == "scope" || s == "err" || s == "op" {
+		return s + "_"
+	}
+	for _, k := range keywords {
+		if s == k {
+			// Alternatively, make the first letter upper case.
+			return s + "_"
+		}
+	}
+	return s
+}
+
+func isListArg(argdef *odpb.OpDef_ArgDef) bool {
+	return argdef.TypeListAttr != "" || argdef.NumberAttr != ""
+}
+
+func isListAttr(attrdef *odpb.OpDef_AttrDef) bool {
+	list, _ := parseTFType(attrdef.Type)
+	return list
+}
+
+func marshalProtoMessage(m proto.Message) string {
+	// Marshal proto message to string.
+	o := prototext.MarshalOptions{Multiline: false}
+	x := o.Format(m)
+
+	// Remove superfluous whitespace, if present.
+	//
+	// Go protobuf output is purposefully unstable, randomly adding
+	// whitespace.  See github.com/golang/protobuf/issues/1121
+	x = strings.ReplaceAll(x, "  ", " ")
+
+	// Remove the prefix of the string up to the first colon.
+	//
+	// This is useful when 's' corresponds to a "oneof" protocol buffer
+	// message. For example, consider the protocol buffer message:
+	//   oneof value { bool b = 1;  int64 i = 2; }
+	// proto.CompactTextString) will print "b:true", or "i:7" etc. The
+	// following strips out the leading "b:" or "i:".
+	y := strings.SplitN(x, ":", 2)
+	if len(y) < 2 {
+		return x
+	}
+	return y[1]
+}
+
+func parseTFType(tfType string) (list bool, typ string) {
+	const (
+		listPrefix = "list("
+		listSuffix = ")"
+	)
+	if strings.HasPrefix(tfType, listPrefix) && strings.HasSuffix(tfType, listSuffix) {
+		return true, strings.TrimSuffix(strings.TrimPrefix(tfType, listPrefix), listSuffix)
+	}
+	return false, tfType
+}

tensorflow/go/genop/internal/genop_test.go

@@ -0,0 +1,824 @@
+/*
+Copyright 2016 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.
+*/
+
+package internal
+
+import (
+	"bytes"
+	"go/format"
+	"testing"
+
+	"google.golang.org/protobuf/encoding/prototext"
+
+	adpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/api_def_go_proto"
+	odpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework/op_def_go_proto"
+)
+
+// Creates an ApiDef based on opdef and applies overrides
+// from apidefText (ApiDef text proto).
+func GetAPIDef(t *testing.T, opdef *odpb.OpDef, apidefText string) *adpb.ApiDef {
+	opdefList := &odpb.OpList{Op: []*odpb.OpDef{opdef}}
+	apimap, err := newAPIDefMap(opdefList)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = apimap.Put(apidefText)
+	if err != nil {
+		t.Fatal(err)
+	}
+	apidef, err := apimap.Get(opdef.Name)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return apidef
+}
+
+// This test is disabled and should be rewritten. Serialized protocol buffers
+// are not stable, see https://github.com/golang/protobuf/issues/1121
+func TestGenerateOp(t *testing.T) {
+
+	// TestGenerateOp validates the generated source code for an op.
+	// The OpDef for the test cases are simplified forms of real ops.
+	testdata := []struct {
+		tag    string
+		opdef  string
+		apidef string
+		wanted string
+	}{
+		{
+			tag: "NoOp",
+			opdef: `
+name: "NoOp"
+`,
+			apidef: `
+op: <
+graph_op_name: "NoOp"
+summary: "No. Op."
+>
+`,
+			wanted: `
+// No. Op.
+//
+// Returns the created operation.
+func NoOp(scope *Scope) (o *tf.Operation) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "NoOp",
+	}
+	return scope.AddOperation(opspec)
+}
+`,
+		},
+		{
+			tag: "NoAttributes",
+			opdef: `
+name: "Add"
+input_arg: <
+  name: "x"
+  type_attr: "T"
+>
+input_arg: <
+  name: "y"
+  type_attr: "T"
+>
+output_arg: <
+  name: "z"
+  type_attr: "T"
+>
+attr: <
+  name: "T"
+  type: "type"
+  allowed_values: <
+    list: <
+      type: DT_FLOAT
+      type: DT_INT64
+    >
+  >
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "Add"
+summary: "Returns x + y element-wise."
+description: "Blah blah",
+>
+`,
+			wanted: `
+// Returns x + y element-wise.
+//
+// Blah blah
+func Add(scope *Scope, x tf.Output, y tf.Output) (z tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	opspec := tf.OpSpec{
+		Type: "Add",
+		Input: []tf.Input{
+			x, y,
+		},
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+`,
+		},
+		{
+			tag: "RequiredAttributes",
+			opdef: `
+name: "Cast"
+input_arg: <
+  name: "x"
+  type_attr: "SrcT"
+>
+output_arg: <
+  name: "y"
+  type_attr: "DstT"
+>
+attr: <
+  name: "SrcT"
+  type: "type"
+>
+attr: <
+  name: "DstT"
+  type: "type"
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "Cast"
+summary: "Cast x of type SrcT to y of DstT."
+>
+`,
+			wanted: `
+// Cast x of type SrcT to y of DstT.
+func Cast(scope *Scope, x tf.Output, DstT tf.DataType) (y tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{"DstT": DstT}
+	opspec := tf.OpSpec{
+		Type: "Cast",
+		Input: []tf.Input{
+			x,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+`,
+		},
+		{
+			tag: "OptionalAttributes",
+			opdef: `
+name: "DecodeJpeg"
+input_arg: <
+  name: "contents"
+  type: DT_STRING
+>
+output_arg: <
+  name: "image"
+  type: DT_UINT8
+>
+attr: <
+  name: "channels"
+  type: "int"
+  default_value: <
+    i: 0
+  >
+>
+attr: <
+  name: "fancy_upscaling"
+  type: "bool"
+  default_value: <
+    b: true
+  >
+>
+attr: <
+  name: "acceptable_fraction"
+  type: "float"
+  default_value: <
+    f: 1
+  >
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "DecodeJpeg"
+in_arg: <
+  name: "contents"
+  description: "0-D.  The JPEG-encoded image."
+>
+out_arg: <
+  name: "image"
+  description: "3-D with shape [height, width, channels]"
+>
+attr: <
+  name: "channels"
+  description: "Number of color channels for the decoded image."
+>
+attr: <
+  name: "fancy_upscaling"
+  description: "If true use a slower but nicer upscaling of the\nchroma planes (yuv420/422 only)."
+>
+attr: <
+  name: "acceptable_fraction"
+  description: "The minimum required fraction of lines before a truncated\ninput is accepted."
+>
+summary: "Decode a JPEG-encoded image to a uint8 tensor."
+description: "Norna dorna fjord\nkajorna\nhahaha"
+>
+`,
+			wanted: `
+// DecodeJpegAttr is an optional argument to DecodeJpeg.
+type DecodeJpegAttr func(optionalAttr)
+
+// DecodeJpegChannels sets the optional channels attribute to value.
+//
+// value: Number of color channels for the decoded image.
+// If not specified, defaults to 0
+func DecodeJpegChannels(value int64) DecodeJpegAttr {
+	return func(m optionalAttr) {
+		m["channels"] = value
+	}
+}
+
+// DecodeJpegFancyUpscaling sets the optional fancy_upscaling attribute to value.
+//
+// value: If true use a slower but nicer upscaling of the
+// chroma planes (yuv420/422 only).
+// If not specified, defaults to true
+func DecodeJpegFancyUpscaling(value bool) DecodeJpegAttr {
+	return func(m optionalAttr) {
+		m["fancy_upscaling"] = value
+	}
+}
+
+// DecodeJpegAcceptableFraction sets the optional acceptable_fraction attribute to value.
+//
+// value: The minimum required fraction of lines before a truncated
+// input is accepted.
+// If not specified, defaults to 1
+func DecodeJpegAcceptableFraction(value float32) DecodeJpegAttr {
+	return func(m optionalAttr) {
+		m["acceptable_fraction"] = value
+	}
+}
+
+// Decode a JPEG-encoded image to a uint8 tensor.
+//
+// Norna dorna fjord
+// kajorna
+// hahaha
+//
+// Arguments:
+//	contents: 0-D.  The JPEG-encoded image.
+//
+// Returns 3-D with shape [height, width, channels]
+func DecodeJpeg(scope *Scope, contents tf.Output, optional ...DecodeJpegAttr) (image tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "DecodeJpeg",
+		Input: []tf.Input{
+			contents,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+`,
+		},
+		{
+			tag: "MultipleOutputs",
+			opdef: `
+name: "TwoOutputs"
+input_arg: <
+  name: "input"
+  type_attr: "T"
+>
+output_arg <
+  name: "x"
+  type_attr: "T"
+>
+output_arg <
+  name: "y"
+  type_attr: "T"
+>
+attr: <
+  name: "T"
+  type: "type"
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "TwoOutputs"
+summary: "Op that produces multiple outputs"
+>
+`,
+			wanted: `
+// Op that produces multiple outputs
+func TwoOutputs(scope *Scope, input tf.Output) (x tf.Output, y tf.Output) {
+        if scope.Err() != nil {
+                return
+        }
+        opspec := tf.OpSpec{
+                Type: "TwoOutputs",
+                Input: []tf.Input{
+                        input,
+                },
+        }
+        op := scope.AddOperation(opspec)
+        return op.Output(0), op.Output(1)
+}
+`,
+		},
+		{
+			tag: "ListOutput",
+			opdef: `
+name: "ShapeN"
+input_arg: <
+  name: "input"
+  type_attr: "T"
+  number_attr: "N"
+>
+output_arg: <
+  name: "output"
+  type_attr: "out_type"
+  number_attr: "N"
+>
+attr: <
+  name: "N"
+  type: "int"
+  has_minimum: true
+  minimum: 1
+>
+attr: <
+  name: "T"
+  type: "type"
+>
+attr: <
+  name: "out_type"
+  type: "type"
+  default_value: <
+    type: DT_INT32
+  >
+  allowed_values: <
+    list: <
+      type: DT_INT32
+      type: DT_INT64
+    >
+  >
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "ShapeN"
+summary: "Returns shape of tensors."
+description: "Some description here."
+>
+`,
+			wanted: `
+// ShapeNAttr is an optional argument to ShapeN.
+type ShapeNAttr func(optionalAttr)
+
+// ShapeNOutType sets the optional out_type attribute to value.
+// If not specified, defaults to DT_INT32
+func ShapeNOutType(value tf.DataType) ShapeNAttr {
+	return func(m optionalAttr) {
+		m["out_type"] = value
+	}
+}
+
+// Returns shape of tensors.
+//
+// Some description here.
+func ShapeN(scope *Scope, input []tf.Output, optional ...ShapeNAttr) (output []tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "ShapeN",
+		Input: []tf.Input{
+			tf.OutputList(input),
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	if scope.Err() != nil {
+		return
+	}
+	var idx int
+	var err error
+	if output, idx, err = makeOutputList(op, idx, "output"); err != nil {
+		scope.UpdateErr("ShapeN", err)
+		return
+	}
+	return output
+}
+`,
+		},
+		{
+			tag: "ApiDefOverrides",
+			opdef: `
+name: "TestOp"
+input_arg: <
+  name: "a"
+  type: DT_STRING
+>
+input_arg: <
+  name: "b"
+  type: DT_STRING
+>
+output_arg: <
+  name: "c"
+  type: DT_UINT8
+>
+attr: <
+  name: "d"
+  type: "int"
+  default_value: <
+    i: 0
+  >
+>
+`,
+			apidef: `
+op: <
+graph_op_name: "TestOp"
+in_arg: <
+  name: "a"
+  rename_to: "aa"
+  description: "Description for aa."
+>
+in_arg: <
+  name: "b"
+  rename_to: "bb"
+  description: "Description for bb."
+>
+arg_order: "b"
+arg_order: "a"
+out_arg: <
+  name: "c"
+  rename_to: "cc"
+  description: "Description for cc."
+>
+attr: <
+  name: "d"
+  rename_to: "dd"
+  description: "Description for dd."
+>
+summary: "Summary for TestOp."
+description: "Description for TestOp."
+>
+`,
+			wanted: `
+// TestOpAttr is an optional argument to TestOp.
+type TestOpAttr func(optionalAttr)
+
+// TestOpDd sets the optional dd attribute to value.
+//
+// value: Description for dd.
+// If not specified, defaults to 0
+func TestOpDd(value int64) TestOpAttr {
+	return func(m optionalAttr) {
+		m["d"] = value
+	}
+}
+
+// Summary for TestOp.
+//
+// Description for TestOp.
+//
+// Arguments:
+//	bb: Description for bb.
+//	aa: Description for aa.
+//
+// Returns Description for cc.
+func TestOp(scope *Scope, bb tf.Output, aa tf.Output, optional ...TestOpAttr) (cc tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "TestOp",
+		Input: []tf.Input{
+			aa, bb,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0)
+}
+`,
+		},
+		{
+			tag: "SampleDistortedBoundingBox",
+			opdef: `
+name: "SampleDistortedBoundingBox"
+input_arg {
+	name: "image_size"
+	type_attr: "T"
+}
+input_arg {
+	name: "bounding_boxes"
+	type: DT_FLOAT
+}
+output_arg {
+	name: "begin"
+	type_attr: "T"
+}
+output_arg {
+	name: "size"
+	type_attr: "T"
+}
+output_arg {
+	name: "bboxes"
+	type: DT_FLOAT
+}
+attr {
+	name: "T"
+	type: "type"
+	allowed_values {
+		list {
+			type: DT_UINT8
+			type: DT_INT8
+			type: DT_INT16
+			type: DT_INT32
+			type: DT_INT64
+		}
+	}
+}
+attr {
+	name: "seed"
+	type: "int"
+	default_value {
+		i: 0
+	}
+}
+attr {
+	name: "seed2"
+	type: "int"
+	default_value {
+		i: 0
+	}
+}
+attr {
+	name: "min_object_covered"
+	type: "float"
+	default_value {
+		f: 0.1
+	}
+}
+attr {
+	name: "aspect_ratio_range"
+	type: "list(float)"
+	default_value {
+		list {
+			f: 0.75
+			f: 1.33
+		}
+	}
+}
+attr {
+	name: "area_range"
+	type: "list(float)"
+	default_value {
+		list {
+			f: 0.05
+			f: 1
+		}
+	}
+}
+attr {
+	name: "max_attempts"
+	type: "int"
+	default_value {
+		i: 100
+	}
+}
+attr {
+	name: "use_image_if_no_bounding_boxes"
+	type: "bool"
+	default_value {
+		b: false
+	}
+}
+is_stateful: true
+`,
+			apidef: `
+op {
+  graph_op_name: "SampleDistortedBoundingBox"
+  in_arg {
+    name: "image_size"
+    description: "Blah blah"
+  }
+  in_arg {
+    name: "bounding_boxes"
+    description: "Blah blah"
+  }
+  out_arg {
+    name: "begin"
+    description: "Blah blah"
+  }
+  out_arg {
+    name: "size"
+    description: "Blah blah"
+  }
+  out_arg {
+    name: "bboxes"
+    description: "Blah blah"
+  }
+  attr {
+    name: "seed"
+    description: "Blah blah"
+  }
+  attr {
+    name: "seed2"
+    description: "Blah blah"
+  }
+  attr {
+    name: "min_object_covered"
+    description: "Blah blah"
+  }
+  attr {
+    name: "aspect_ratio_range"
+    description: "Blah blah"
+  }
+  attr {
+    name: "area_range"
+    description: "Blah blah"
+  }
+  attr {
+    name: "max_attempts"
+    description: "Blah blah"
+  }
+  attr {
+    name: "use_image_if_no_bounding_boxes"
+    description: "Blah blah"
+  }
+  summary: "Generate a single randomly distorted bounding box for an image."
+	description: "Blah blah"
+}
+`,
+			wanted: `
+// SampleDistortedBoundingBoxAttr is an optional argument to SampleDistortedBoundingBox.
+type SampleDistortedBoundingBoxAttr func(optionalAttr)
+
+// SampleDistortedBoundingBoxSeed sets the optional seed attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to 0
+func SampleDistortedBoundingBoxSeed(value int64) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["seed"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxSeed2 sets the optional seed2 attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to 0
+func SampleDistortedBoundingBoxSeed2(value int64) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["seed2"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxMinObjectCovered sets the optional min_object_covered attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to 0.1
+func SampleDistortedBoundingBoxMinObjectCovered(value float32) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["min_object_covered"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxAspectRatioRange sets the optional aspect_ratio_range attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to {f:0.75 f:1.33}
+func SampleDistortedBoundingBoxAspectRatioRange(value []float32) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["aspect_ratio_range"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxAreaRange sets the optional area_range attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to {f:0.05 f:1}
+func SampleDistortedBoundingBoxAreaRange(value []float32) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["area_range"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxMaxAttempts sets the optional max_attempts attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to 100
+func SampleDistortedBoundingBoxMaxAttempts(value int64) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["max_attempts"] = value
+	}
+}
+
+// SampleDistortedBoundingBoxUseImageIfNoBoundingBoxes sets the optional use_image_if_no_bounding_boxes attribute to value.
+//
+// value: Blah blah
+// If not specified, defaults to false
+func SampleDistortedBoundingBoxUseImageIfNoBoundingBoxes(value bool) SampleDistortedBoundingBoxAttr {
+	return func(m optionalAttr) {
+		m["use_image_if_no_bounding_boxes"] = value
+	}
+}
+
+// Generate a single randomly distorted bounding box for an image.
+//
+// Blah blah
+//
+// Arguments:
+//	image_size: Blah blah
+//	bounding_boxes: Blah blah
+//
+// Returns:
+//	begin: Blah blah
+//	size: Blah blah
+//	bboxes: Blah blah
+func SampleDistortedBoundingBox(scope *Scope, image_size tf.Output, bounding_boxes tf.Output, optional ...SampleDistortedBoundingBoxAttr) (begin tf.Output, size tf.Output, bboxes tf.Output) {
+	if scope.Err() != nil {
+		return
+	}
+	attrs := map[string]interface{}{}
+	for _, a := range optional {
+		a(attrs)
+	}
+	opspec := tf.OpSpec{
+		Type: "SampleDistortedBoundingBox",
+		Input: []tf.Input{
+			image_size, bounding_boxes,
+		},
+		Attrs: attrs,
+	}
+	op := scope.AddOperation(opspec)
+	return op.Output(0), op.Output(1), op.Output(2)
+}
+`,
+		},
+	}
+
+	for _, test := range testdata {
+		t.Run(test.tag, func(t *testing.T) {
+			var opdef odpb.OpDef
+			var apidef *adpb.ApiDef
+			var buf bytes.Buffer
+			if err := prototext.Unmarshal([]byte(test.opdef), &opdef); err != nil {
+				t.Fatal(err)
+			}
+			apidef = GetAPIDef(t, &opdef, test.apidef)
+			if err := generateFunctionForOp(&buf, &opdef, apidef); err != nil {
+				t.Fatal(err)
+			}
+			got, err := format.Source(buf.Bytes())
+			if err != nil {
+				t.Fatalf("Unable to format: %v\n%s", err, buf.Bytes())
+			}
+			want, err := format.Source([]byte(test.wanted))
+			if err != nil {
+				t.Fatalf("Unable to format: %v\n%s", err, test.wanted)
+			}
+			if !bytes.Equal(got, want) {
+				t.Fatalf("Got:\n%s\nWant:\n%s\n", got, want)
+			}
+		})
+	}
+}

tensorflow/go/genop/internal/lib.go

@@ -0,0 +1,22 @@
+/*
+Copyright 2016 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.
+*/
+
+// Package internal generates Go source code with functions for TensorFlow operations.
+package internal
+
+// #cgo LDFLAGS: -ltensorflow
+// #cgo CFLAGS: -I${SRCDIR}/../../../../
+import "C"

tensorflow/go/genop/main.go

@@ -0,0 +1,69 @@
+/*
+Copyright 2016 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.
+*/
+
+// Command genop generates a Go source file with functions for TensorFlow ops.
+package main
+
+import (
+	"bytes"
+	"flag"
+	"go/format"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/tensorflow/tensorflow/tensorflow/go/genop/internal"
+)
+
+func main() {
+	var (
+		filename   = flag.String("outfile", "", "File to write generated source code to.")
+		header     = flag.String("header", "", "Path to a file whose contents will be copied into the generated file. Can be empty")
+		apiDefDirs = flag.String("api_def_dirs", "", "Comma-separated directories containing api_def_*.pbtxt files.")
+		buf        bytes.Buffer
+	)
+	flag.Parse()
+	if *filename == "" {
+		log.Fatal("-outfile must be set")
+	}
+	if *header != "" {
+		hdr, err := os.ReadFile(*header)
+		if err != nil {
+			log.Fatalf("Unable to read %s: %v", *header, err)
+		}
+		buf.Write(hdr)
+		buf.WriteString("\n\n")
+	}
+	os.MkdirAll(filepath.Dir(*filename), 0755)
+
+	apiDefDirsList := []string{}
+	if len(*apiDefDirs) > 0 {
+		apiDefDirsList = strings.Split(*apiDefDirs, ",")
+	}
+
+	if err := internal.GenerateFunctionsForRegisteredOps(
+		&buf, apiDefDirsList); err != nil {
+		log.Fatal(err)
+	}
+	formatted, err := format.Source(buf.Bytes())
+	if err != nil {
+		log.Fatalf("Failed to generate valid source? 'go fmt' failed: %v", err)
+	}
+	if err := os.WriteFile(*filename, formatted, 0644); err != nil {
+		log.Fatalf("Failed to write to %q: %v", *filename, err)
+	}
+}