def where(condition, x=pu.dummy, y=pu.dummy): a = condition s = arrayfire.where(a.d_array) # numpy uses int64 while arrayfire uses uint32 s = ndarray(pu.af_shape(s), dtype=numpy.uint32, af_array=s).astype(numpy.int64) if(x is pu.dummy and y is pu.dummy): idx = [] mult = 1 for i in a.shape[::-1]: mult *= i idx = [s % mult] + idx s /= mult idx = tuple(idx) return idx elif(x is not pu.dummy and y is not pu.dummy): if(x.dtype != y.dtype): raise TypeError('x and y must have same dtype') if(x.shape != y.shape): raise ValueError('x and y must have same shape') ret = array(y) if(len(ret.shape) > 1): ret = ret.flatten() ret[s] = x.flatten()[s] ret = ret.reshape(x.shape) else: ret[s] = x[s] return ret; else: raise ValueError('either both or neither of x and y should be given')
def transpose(self, *axes): if(self.ndim == 1): return self if len(axes) == 0 and self.ndim == 2: s = arrayfire.transpose(self.d_array) else: order = [0,1,2,3] if len(axes) == 0 or axes[0] is None: order[:self.ndim] = order[:self.ndim][::-1] else: if isinstance(axes[0], collections.Iterable): axes = axes[0] for i,ax in enumerate(axes): order[i] = pu.c2f(self.shape, ax) # We have to do this gymnastic due to the fact that arrayfire # uses Fortran order order[:len(axes)] = order[:len(axes)][::-1] #print order s = arrayfire.reorder(self.d_array, order[0],order[1],order[2],order[3]) return ndarray(pu.af_shape(s), dtype=self.dtype, af_array=s)
def __getitem__(self, args): if not isinstance(args, tuple): args = (args,) if len(args) == 1 and isinstance(args[0], afnumpy.ndarray) and args[0].dtype == numpy.dtype('bool'): # Special case for boolean getitem return self.flat[afnumpy.where(args[0].flat)] idx, new_shape = indexing.__convert_dim__(self.shape, args) if any(x is None for x in idx): # one of the indices is empty return ndarray(indexing.__index_shape__(self.shape, idx), dtype=self.dtype) idx = tuple(idx) if len(idx) == 0: idx = tuple([0]) s = self.d_array[idx] shape = pu.af_shape(s) array = ndarray(shape, dtype=self.dtype, af_array=s) if(shape != new_shape): array = array.reshape(new_shape) if new_shape == () and Ellipsis not in args: # Return the actual scalar return numpy.array(array)[()] return array