tensor_int32.go

// This file was automatically generated by genny.
// Any changes will be lost if this file is regenerated.
// see https://github.com/cheekybits/genny

package tensorflow

import (
	"fmt"
	"unsafe"
)

// #include <string.h>
// #include <tensor_c_api.h>
// #include "tensorflow.h"
import "C"

var _ Tensor = &Int32Tensor{}
var _ Tensor = &Float32Tensor{}

type Int32Tensor struct {
	dims []int
	data []int32
}

func NewInt32Tensor(dims []int) *Int32Tensor {
	size := 1
	for _, dim := range dims {
		size *= dim
	}

	return &Int32Tensor{
		dims: dims,
		data: make([]int32, size),
	}
}

func (t *Int32Tensor) index(idx []int) (int, int) {
	if len(idx) >= len(t.dims) {
		panic(fmt.Sprintf("Trying to address using %d dimensions, only %d permitted",
			len(idx), len(t.dims)-1))
	}

	// Special case: the full array
	if len(idx) == 0 {
		dimSize := 1
		for _, dim := range t.dims {
			dimSize *= dim
		}

		return 0, dimSize
	}

	startIdx := 0
	var dimSize int
	for i, idx := range idx {
		dimSize = 1
		for _, dim := range t.dims[i+1:] {
			dimSize *= dim
		}

		startIdx += idx * dimSize
	}

	return startIdx, dimSize
}

func (t *Int32Tensor) Get(idx []int) []int32 {
	startIdx, dimSize := t.index(idx)

	return t.data[startIdx : startIdx+dimSize]
}

func (t *Int32Tensor) Assign(idx []int, data []int32) {
	startIdx, _ := t.index(idx)

	if len(idx) >= len(t.dims) {
		panic(fmt.Sprintf("Trying to address using %d dimensions, only %d permitted",
			len(idx), len(t.dims)-1))
	}

	// TODO: bounds checks?

	copy(t.data[startIdx:], data)
}

func (t *Int32Tensor) Fill(idx []int, v int32) {
	startIdx, dimSize := t.index(idx)

	for idx := startIdx; idx < startIdx+dimSize; idx++ {
		t.data[idx] = v
	}
}

func (t *Int32Tensor) toCTensor() *C.TF_Tensor {
	// TF_NewTensor copies dims, does not take ownership.
	llDims := make([]C.longlong, len(t.dims))
	for idx, val := range t.dims {
		llDims[idx] = C.longlong(val)
	}

	// TF_NewTensor potentially adopts data, so we cannot pass the Go
	// slice backing array.
	dataLen := C.size_t(len(t.data)) * C.size_t(unsafe.Sizeof(t.data[0]))
	cData := C.malloc(dataLen)
	C.memcpy(cData, unsafe.Pointer(&t.data[0]), dataLen)

	return C.tfgo_tensor(C.TF_INT32, (*C.longlong)(unsafe.Pointer(&llDims[0])),
		C.int(len(llDims)), cData, dataLen)
}

func adoptint32Tensor(ct *C.TF_Tensor) *Int32Tensor {
	dims := C.TF_NumDims(ct)
	shape := make([]int, dims)
	size := uint(1)
	for i := C.int(0); i < dims; i++ {
		shape[i] = int(C.TF_Dim(ct, i))
		size *= uint(shape[i])
	}

	bs := C.TF_TensorByteSize(ct)
	var valForSize int32
	if uint(bs)/uint(unsafe.Sizeof(valForSize)) != size {
		panic("Expected tensor size does not correspond to the actual tensor size")
	}

	data := make([]int32, size)
	C.memcpy(unsafe.Pointer(&data[0]), C.TF_TensorData(ct), bs)

	return &Int32Tensor{
		dims: shape,
		data: data,
	}
}