예제 #1
0
func newOpBuilder(g *Graph, typ string, name string) *opBuilder {
	opType := C.CString(typ)
	opName := C.CString(name)
	b := &opBuilder{c: C.TF_NewOperation(g.c, opType, opName), g: g}
	C.free(unsafe.Pointer(opType))
	C.free(unsafe.Pointer(opName))
	return b
}
예제 #2
0
// AddOperation adds an operation to g.
func (g *Graph) AddOperation(args OpSpec) (*Operation, error) {
	if args.Name == "" {
		args.Name = args.Type
	}
	cname := C.CString(args.Name)
	ctype := C.CString(args.Type)
	cdesc := C.TF_NewOperation(g.c, ctype, cname)
	C.free(unsafe.Pointer(cname))
	C.free(unsafe.Pointer(ctype))

	for _, in := range args.Input {
		switch in := in.(type) {
		case Output:
			C.TF_AddInput(cdesc, in.c())
		case OutputList:
			size := len(in)
			list := make([]C.TF_Output, size)
			for i, v := range in {
				list[i] = v.c()
			}
			C.TF_AddInputList(cdesc, &list[0], C.int(size))
		}
	}
	status := newStatus()
	for name, value := range args.Attrs {
		if err := setAttr(cdesc, status, name, value); err != nil {
			// Memory leak here as the TF_OperationDescription
			// object will not be cleaned up. At the time of this
			// writing, this was next to impossible since it
			// required value to be a string tensor with
			// incorrectly encoded strings. Given this rarity, live
			// with the memory leak.  If it becomes a real problem,
			// consider adding a TF_DeleteOperationDescription
			// function to the C API.
			return nil, fmt.Errorf("%v (memory will be leaked)", err)
		}
	}
	op := &Operation{
		c: C.TF_FinishOperation(cdesc, status.c),
		g: g,
	}
	return op, status.Err()
}