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
