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)}