lib/tensorflow/tensor.rb in tensorflow-0.1.0 vs lib/tensorflow/tensor.rb in tensorflow-0.1.1
- old
+ new
@@ -1,8 +1,8 @@
module TensorFlow
class Tensor
- def initialize(value = nil, pointer: nil, dtype: nil, shape: nil)
+ def initialize(value = nil, dtype: nil, shape: nil, pointer: nil)
@status = FFI.TF_NewStatus
if pointer
@pointer = pointer
else
@@ -16,21 +16,31 @@
dims_ptr = nil
end
data = data.flatten
- dtype ||= Utils.infer_type(value)
+ dtype ||= Utils.infer_type(data)
type = FFI::DataType[dtype]
case dtype
+ when :float, :double, :int32, :uint8, :int16, :int8, :int64, :uint16, :uint32, :uint64
+ data_ptr = ::FFI::MemoryPointer.new(dtype, data.size)
+ data_ptr.send("write_array_of_#{dtype}", data)
+ when :bfloat16
+ # https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
+ data_ptr = ::FFI::MemoryPointer.new(:int8, data.size * 2)
+ data_ptr.write_bytes(data.map { |v| [v].pack("g")[0..1] }.join)
+ when :complex64
+ data_ptr = ::FFI::MemoryPointer.new(:float, data.size * 2)
+ data_ptr.write_array_of_float(data.flat_map { |v| [v.real, v.imaginary] })
+ when :complex128
+ data_ptr = ::FFI::MemoryPointer.new(:double, data.size * 2)
+ data_ptr.write_array_of_double(data.flat_map { |v| [v.real, v.imaginary] })
when :string
data_ptr = string_ptr(data)
- when :float
- data_ptr = ::FFI::MemoryPointer.new(:float, data.size)
- data_ptr.write_array_of_float(data)
- when :int32
- data_ptr = ::FFI::MemoryPointer.new(:int32, data.size)
- data_ptr.write_array_of_int32(data)
+ when :bool
+ data_ptr = ::FFI::MemoryPointer.new(:int8, data.size)
+ data_ptr.write_array_of_int8(data.map { |v| v ? 1 : 0 })
else
raise "Unknown type: #{dtype}"
end
callback = ::FFI::Function.new(:void, [:pointer, :size_t, :pointer]) do |data, len, arg|
@@ -40,78 +50,45 @@
tensor = FFI.TF_NewTensor(type, dims_ptr, shape.size, data_ptr, data_ptr.size, callback, nil)
@pointer = FFI.TFE_NewTensorHandle(tensor, @status)
check_status @status
end
- # TODO fix segfault
- # ObjectSpace.define_finalizer(self, self.class.finalize(@pointer))
+ ObjectSpace.define_finalizer(self, self.class.finalize(@pointer, @status, tensor))
end
def +(other)
- TensorFlow.add(self, other)
+ Math.add(self, other)
end
def -(other)
- TensorFlow.subtract(self, other)
+ Math.subtract(self, other)
end
def *(other)
- TensorFlow.multiply(self, other)
+ Math.multiply(self, other)
end
def /(other)
- TensorFlow.divide(self, other)
+ Math.divide(self, other)
end
def %(other)
- TensorFlow.floormod(self, other)
+ Math.floormod(self, other)
end
- def num_dims
- ret = FFI.TFE_TensorHandleNumDims(@pointer, @status)
- check_status @status
- ret
- end
-
- def dtype
- @dtype ||= FFI::DataType[FFI.TFE_TensorHandleDataType(@pointer)]
- end
-
- def element_count
- ret = FFI.TFE_TensorHandleNumElements(@pointer, @status)
- check_status @status
- ret
- end
-
- def shape
- @shape ||= begin
- shape = []
- num_dims.times do |i|
- shape << FFI.TFE_TensorHandleDim(@pointer, i, @status)
- check_status @status
- end
- shape
- end
- end
-
- def data_pointer
- tensor = FFI.TFE_TensorHandleResolve(@pointer, @status)
- check_status @status
- FFI.TF_TensorData(tensor)
- end
-
- def to_ptr
- @pointer
- end
-
def value
value =
case dtype
- when :float
- data_pointer.read_array_of_float(element_count)
- when :int32
- data_pointer.read_array_of_int32(element_count)
+ when :float, :double, :int32, :uint8, :int16, :int8, :int64, :uint16, :uint32, :uint64
+ data_pointer.send("read_array_of_#{dtype}", element_count)
+ when :bfloat16
+ byte_str = data_pointer.read_bytes(element_count * 2)
+ element_count.times.map { |i| "#{byte_str[(2 * i)..(2 * i + 1)]}\x00\x00".unpack1("g") }
+ when :complex64
+ data_pointer.read_array_of_float(element_count * 2).each_slice(2).map { |v| Complex(*v) }
+ when :complex128
+ data_pointer.read_array_of_double(element_count * 2).each_slice(2).map { |v| Complex(*v) }
when :string
# string tensor format
# https://github.com/tensorflow/tensorflow/blob/5453aee48858fd375172d7ae22fad1557e8557d6/tensorflow/c/tf_tensor.h#L57
start_offset_size = element_count * 8
offsets = data_pointer.read_array_of_uint64(element_count)
@@ -125,10 +102,25 @@
end
reshape(value, shape)
end
+ def dtype
+ @dtype ||= FFI::DataType[FFI.TFE_TensorHandleDataType(@pointer)]
+ end
+
+ def shape
+ @shape ||= begin
+ shape = []
+ num_dims.times do |i|
+ shape << FFI.TFE_TensorHandleDim(@pointer, i, @status)
+ check_status @status
+ end
+ shape
+ end
+ end
+
def to_s
inspect
end
def to_i
@@ -137,22 +129,48 @@
def to_a
value
end
+ def to_ptr
+ @pointer
+ end
+
def inspect
inspection = %w(value shape dtype).map { |v| "#{v}: #{send(v).inspect}"}
"#<#{self.class} #{inspection.join(", ")}>"
end
- def self.finalize(pointer)
+ def self.finalize(pointer, status, tensor)
# must use proc instead of stabby lambda
- proc { FFI.TFE_DeleteTensorHandle(pointer) }
+ proc do
+ FFI.TFE_DeleteTensorHandle(pointer)
+ FFI.TFE_DeleteStatus(status)
+ FFI.TFE_DeleteTensor(tensor) if tensor
+ end
end
private
+ def num_dims
+ ret = FFI.TFE_TensorHandleNumDims(@pointer, @status)
+ check_status @status
+ ret
+ end
+
+ def element_count
+ ret = FFI.TFE_TensorHandleNumElements(@pointer, @status)
+ check_status @status
+ ret
+ end
+
+ def data_pointer
+ tensor = FFI.TFE_TensorHandleResolve(@pointer, @status)
+ check_status @status
+ FFI.TF_TensorData(tensor)
+ end
+
def reshape(arr, dims)
return arr.first if dims.empty?
arr = arr.flatten
dims[1..-1].reverse.each do |dim|
arr = arr.each_slice(dim)
@@ -174,15 +192,15 @@
# https://github.com/tensorflow/tensorflow/blob/5453aee48858fd375172d7ae22fad1557e8557d6/tensorflow/c/tf_tensor.h#L57
def string_ptr(data)
start_offset_size = data.size * 8
offsets = [0]
data.each do |str|
- offsets << str.bytesize + 1
+ offsets << offsets.last + str.bytesize + 1
end
data_ptr = ::FFI::MemoryPointer.new(:char, start_offset_size + offsets.pop)
data_ptr.write_array_of_uint64(offsets)
- data.zip(offsets) do |str, off|
- (data_ptr + start_offset_size + off).write_string(str)
+ data.zip(offsets) do |str, offset|
+ (data_ptr + start_offset_size + offset).write_string(str)
end
data_ptr
end
def check_status(status)