def __init__(self, dtype=None, sizes=[], strides=None, ptr=None):
if ptr:
ffi_obj = ptr
elif dtype is not None:
if strides is None:
strides = []
if len(strides) != len(sizes):
stride = 1
for i in range(len(sizes) - 1, -1, -1):
strides.insert(0, stride)
stride *= sizes[i]
raw_sizes = ffi.new('int64_t[]', sizes)
raw_strides = ffi.new('int64_t[]', strides)
ffi_obj = ffi_call(lib.plaidml_shape_alloc, dtype, len(sizes), raw_sizes, raw_strides)
else:
raise ValueError('One of dtype= or ptr= must be specified.')
super(TensorShape, self).__init__(ffi_obj)
def __init__(self, dtype=None, dims=[], ptr=None):
if ptr:
ffi_obj = ptr
elif dtype is not None:
raw_dims = ffi.new('int64_t[]', [0 if x is None else x for x in dims])
ffi_obj = ffi_call(lib.plaidml_logical_shape_alloc, dtype, len(dims), raw_dims)
else:
raise ValueError('One of dtype= or ptr= must be specified.')
super(LogicalShape, self).__init__(ffi_obj)
def gradients(loss, variables):
wrts = [x.as_ptr() for x in variables]
raw_grads = ffi.new('plaidml_expr*[]', len(wrts))
ffi_call(
lib.plaidml_expr_gradient,
len(wrts),
wrts,
loss.as_ptr(),
raw_grads,
)
return [Tensor(expr=x) for x in raw_grads]
def __init__(self, name, outputs, updates=[]):
raw_outputs = [x.as_ptr() for x in outputs]
new_outputs = ffi.new('plaidml_expr*[]', len(outputs))
dst_updates = [x[0].as_ptr() for x in updates]
src_updates = [x[1].as_ptr() for x in updates]
ffi_obj = ffi_call(
lib.plaidml_program_evaluate,
name.encode(),
len(raw_outputs),
raw_outputs,
new_outputs,
len(updates),
src_updates,
dst_updates,
)
self.outputs = [Tensor(expr=x) for x in new_outputs]
super(Program, self).__init__(ffi_obj)
def __init__(self, name, outputs, updates=[]):
raw_outputs = [x.as_ptr() for x in outputs]
dst_updates = [x[0].as_ptr() for x in updates]
src_updates = [x[1].as_ptr() for x in updates]
raw_args = ffi.new('plaidml_program_args**')
ffi_obj = ffi_call(
lib.plaidml_program_evaluate,
name.encode(),
len(raw_outputs),
raw_outputs,
len(updates),
src_updates,
dst_updates,
raw_args,
)
self.args = [
ProgramArgument(raw_args[0].args[i])
for i in range(raw_args[0].nargs)
]
ffi_call(lib.plaidml_program_args_free, raw_args[0])
super(Program, self).__init__(ffi_obj)
def wrap(x, y):
return ffi.new('plaidml_binding*', [x.as_ptr(), y.as_ptr()])
def list_targets():
ntargets = ffi_call(lib.plaidml_target_list_count)
raw_targets = ffi.new('plaidml_string*[]', ntargets)
ffi_call(lib.plaidml_target_list, ntargets, raw_targets)
return [decode_str(x) for x in raw_targets]
def list_devices():
ndevices = ffi_call(lib.plaidml_device_list_count)
raw_devices = ffi.new('plaidml_string*[]', ndevices)
ffi_call(lib.plaidml_device_list, ndevices, raw_devices)
return [decode_str(x) for x in raw_devices]
def all():
nitems = ffi_call(lib.plaidml_settings_list_count)
raw_keys = ffi.new('plaidml_string*[]', nitems)
raw_values = ffi.new('plaidml_string*[]', nitems)
ffi_call(lib.plaidml_settings_list, nitems, raw_keys, raw_values)
return {decode_str(key): decode_str(value) for key, value in zip(raw_keys, raw_values)}
