go: Add Operation.OutputListSize

This will be needed for generating the function wrappers for ops that produce a
list of tensors as output.

Another step towards #10
Change: 136191993

1 files changed, 16 insertions, 0 deletions

diff --git a/tensorflow/go/operation.go b/tensorflow/go/operation.go
index 8f4fee1cbe..e7986814d0 100644
--- a/tensorflow/go/operation.go
+++ b/tensorflow/go/operation.go
@@ -20,6 +20,7 @@ import "C"
 
 import (
 	"errors"
+	"unsafe"
 )
 
 // Operation that has been added to the graph.
@@ -45,6 +46,21 @@ func (op *Operation) NumOutputs() int {
 	return int(C.TF_OperationNumOutputs(op.c))
 }
 
+// OutputListSize returns the size of the list of Outputs that is produced by a
+// named output of op.
+//
+// An Operation has multiple named outputs, each of which produces either
+// a single tensor or a list of tensors. This method returns the size of
+// the list of tensors for a specific output of the operation, identified
+// by its name.
+func (op *Operation) OutputListSize(output string) (int, error) {
+	cname := C.CString(output)
+	defer C.free(unsafe.Pointer(cname))
+	status := newStatus()
+	n := C.TF_OperationOutputListLength(op.c, cname, status.c)
+	return int(n), status.Err()
+}
+
 // Output returns the i-th output of op.
 func (op *Operation) Output(i int) Output {
 	return Output{op, i}