func (b *opBuilder) SetAttrType(name string, typ DataType) {
attrName := C.CString(name)
C.TF_SetAttrType(b.c, attrName, C.TF_DataType(typ))
C.free(unsafe.Pointer(attrName))
}
func setAttr(cdesc *C.TF_OperationDescription, status *status, name string, value interface{}) error {
cAttrName := C.CString(name)
defer C.free(unsafe.Pointer(cAttrName))
switch value := value.(type) {
case string:
cstr := C.CString(value)
C.TF_SetAttrString(cdesc, cAttrName, unsafe.Pointer(cstr), C.size_t(len(value)))
C.free(unsafe.Pointer(cstr))
case []string:
size := len(value)
list := make([]unsafe.Pointer, size)
lens := make([]C.size_t, size)
for i, s := range value {
list[i] = unsafe.Pointer(C.CString(s))
lens[i] = C.size_t(len(s))
}
C.TF_SetAttrStringList(cdesc, cAttrName, &list[0], &lens[0], C.int(size))
for _, s := range list {
C.free(s)
}
case int64:
C.TF_SetAttrInt(cdesc, cAttrName, C.int64_t(value))
case []int64:
size := len(value)
list := make([]C.int64_t, size)
for i, v := range value {
list[i] = C.int64_t(v)
}
C.TF_SetAttrIntList(cdesc, cAttrName, &list[0], C.int(size))
case float32:
C.TF_SetAttrFloat(cdesc, cAttrName, C.float(value))
case []float32:
size := len(value)
list := make([]C.float, size)
for i, v := range value {
list[i] = C.float(v)
}
C.TF_SetAttrFloatList(cdesc, cAttrName, &list[0], C.int(size))
case bool:
v := C.uchar(0)
if value {
v = 1
}
C.TF_SetAttrBool(cdesc, cAttrName, v)
case []bool:
size := len(value)
list := make([]C.uchar, size)
for i, v := range value {
if v {
list[i] = 1
}
}
C.TF_SetAttrBoolList(cdesc, cAttrName, &list[0], C.int(size))
case DataType:
C.TF_SetAttrType(cdesc, cAttrName, C.TF_DataType(value))
case []DataType:
list := (*C.TF_DataType)(&value[0])
C.TF_SetAttrTypeList(cdesc, cAttrName, list, C.int(len(value)))
case *Tensor:
C.TF_SetAttrTensor(cdesc, cAttrName, value.c, status.c)
if err := status.Err(); err != nil {
return fmt.Errorf("bad value for attribute %q: %v", name, err)
}
case []*Tensor:
size := len(value)
list := make([]*C.TF_Tensor, size)
for i, v := range value {
list[i] = v.c
}
C.TF_SetAttrTensorList(cdesc, cAttrName, &list[0], C.int(size), status.c)
if err := status.Err(); err != nil {
return fmt.Errorf("bad value for attribute %q: %v", name, err)
}
case Shape:
ndims, dims := cshape(value)
var dimsp *C.int64_t
if ndims > 0 {
dimsp = &dims[0]
}
C.TF_SetAttrShape(cdesc, cAttrName, dimsp, ndims)
case []Shape:
ndims := make([]C.int, len(value))
dims := make([][]C.int64_t, len(value))
dimsp := make([]*C.int64_t, len(value))
for i, s := range value {
ndims[i], dims[i] = cshape(s)
if ndims[i] > 0 {
dimsp[i] = &dims[i][0]
}
}
C.TF_SetAttrShapeList(cdesc, cAttrName, &dimsp[0], &ndims[0], C.int(len(value)))
default:
return fmt.Errorf("attribute %q has a type (%T) which is not valid for operation attributes", name, value)
}
return nil
}
func setAttr(cdesc *C.TF_OperationDescription, status *status, name string, value interface{}) error {
cAttrName := C.CString(name)
defer C.free(unsafe.Pointer(cAttrName))
switch value := value.(type) {
case string:
cstr := C.CString(value)
C.TF_SetAttrString(cdesc, cAttrName, unsafe.Pointer(cstr), C.size_t(len(value)))
C.free(unsafe.Pointer(cstr))
case []string:
size := len(value)
list := make([]unsafe.Pointer, size)
lens := make([]C.size_t, size)
for i, s := range value {
list[i] = unsafe.Pointer(C.CString(s))
lens[i] = C.size_t(len(s))
}
C.TF_SetAttrStringList(cdesc, cAttrName, &list[0], &lens[0], C.int(size))
for _, s := range list {
C.free(s)
}
case int64:
C.TF_SetAttrInt(cdesc, cAttrName, C.int64_t(value))
case []int64:
size := len(value)
list := make([]C.int64_t, size)
for i, v := range value {
list[i] = C.int64_t(v)
}
C.TF_SetAttrIntList(cdesc, cAttrName, &list[0], C.int(size))
case float32:
C.TF_SetAttrFloat(cdesc, cAttrName, C.float(value))
case []float32:
size := len(value)
list := make([]C.float, size)
for i, v := range value {
list[i] = C.float(v)
}
C.TF_SetAttrFloatList(cdesc, cAttrName, &list[0], C.int(size))
case bool:
v := C.uchar(0)
if value {
v = 1
}
C.TF_SetAttrBool(cdesc, cAttrName, v)
case []bool:
size := len(value)
list := make([]C.uchar, size)
for i, v := range value {
if v {
list[i] = 1
}
}
C.TF_SetAttrBoolList(cdesc, cAttrName, &list[0], C.int(size))
case DataType:
C.TF_SetAttrType(cdesc, cAttrName, C.TF_DataType(value))
case []DataType:
list := (*C.TF_DataType)(&value[0])
C.TF_SetAttrTypeList(cdesc, cAttrName, list, C.int(len(value)))
case *Tensor:
C.TF_SetAttrTensor(cdesc, cAttrName, value.c, status.c)
if err := status.Err(); err != nil {
return fmt.Errorf("bad value for attribute %q: %v", name, err)
}
case []*Tensor:
size := len(value)
list := make([]*C.TF_Tensor, size)
for i, v := range value {
list[i] = v.c
}
C.TF_SetAttrTensorList(cdesc, cAttrName, &list[0], C.int(size), status.c)
if err := status.Err(); err != nil {
return fmt.Errorf("bad value for attribute %q: %v", name, err)
}
default:
// Shapes can be done, but will require that it be
// distinguishable from []int64. Which is fine, it
// probably makes sense to define a Shape type anyway,
// since that should handle partially known shapes as
// well and hide the special meaning of -1?
return fmt.Errorf("attribute %q has a type (%T) which is not valid for operation attributes", name, value)
}
return nil
}