def __index_shape__(A_shape, idx, del_singleton=True):
shape = []
for i in range(0,len(idx)):
if(idx[i] is None):
shape.append(0)
elif(isinstance(idx[i],numbers.Number)):
if del_singleton:
# Remove dimensions indexed with a scalar
continue
else:
shape.append(1)
elif(isinstance(idx[i],arrayfire.index.Seq)):
if(idx[i].s == arrayfire.af_span):
shape.append(A_shape[i])
else:
shape.append(idx[i].size)
elif(isinstance(idx[i],slice)):
shape.append(__slice_len__(idx[i], pu.c2f(A_shape), i))
elif(isinstance(idx[i], arrayfire.Array)):
shape.append(idx[i].elements())
elif(isinstance(idx[i],arrayfire.index)):
if(idx[i].isspan()):
shape.append(A_shape[i])
else:
af_idx = idx[i].get()
if(af_idx.isBatch):
raise ValueError
if(af_idx.isSeq):
shape.append(arrayfire.seq(af_idx.seq()).size)
else:
shape.append(af_idx.arr_elements())
else:
raise ValueError
return pu.c2f(shape)
def __slice_to_seq__(shape, idx, axis):
maxlen = shape[axis]
if(isinstance(idx, numbers.Number)):
if idx < 0:
idx = maxlen + idx
if(idx >= maxlen):
raise IndexError('index %d is out of bounds for axis %d with size %d' % (idx, axis, maxlen))
return idx
if(isinstance(idx, afnumpy.ndarray)):
return idx.d_array
if not isinstance(idx, slice):
return afnumpy.array(idx).d_array
if idx.step is None:
step = 1
else:
step = idx.step
if idx.start is None:
if step < 0:
start = maxlen-1
else:
start = 0
else:
start = idx.start
if(start < 0):
start += maxlen
if idx.stop is None:
if step < 0:
end = 0
else:
end = maxlen-1
else:
end = idx.stop
if(end < 0):
end += maxlen
if step < 0:
end += 1
else:
end -= 1
# arrayfire doesn't like other steps in this case
if(start == end):
step = 1
if((start-end > 0 and step > 0) or
(start-end < 0 and step < 0)):
return None
return arrayfire.seq(float(start),
float(end),
float(step))
def __slice_to_seq__(shape, idx, axis):
maxlen = shape[axis]
if (isinstance(idx, numbers.Number)):
if idx < 0:
idx = maxlen + idx
if (idx >= maxlen):
raise IndexError(
'index %d is out of bounds for axis %d with size %d' %
(idx, axis, maxlen))
return idx
if (isinstance(idx, afnumpy.ndarray)):
return idx.d_array
if not isinstance(idx, slice):
return afnumpy.array(idx).d_array
if idx.step is None:
step = 1
else:
step = idx.step
if idx.start is None:
if step < 0:
start = maxlen - 1
else:
start = 0
else:
start = idx.start
if (start < 0):
start += maxlen
if idx.stop is None:
if step < 0:
end = 0
else:
end = maxlen - 1
else:
end = idx.stop
if (end < 0):
end += maxlen
if step < 0:
end += 1
else:
end -= 1
# arrayfire doesn't like other steps in this case
if (start == end):
step = 1
if ((start - end > 0 and step > 0) or (start - end < 0 and step < 0)):
return None
return arrayfire.seq(float(start), float(end), float(step))
def __index_shape__(A_shape, idx, del_singleton=True):
shape = []
for i in range(0, len(idx)):
if (idx[i] is None):
shape.append(0)
elif (isinstance(idx[i], numbers.Number)):
if del_singleton:
# Remove dimensions indexed with a scalar
continue
else:
shape.append(1)
elif (isinstance(idx[i], arrayfire.index.Seq)):
if (idx[i].s == arrayfire.af_span):
shape.append(A_shape[i])
else:
shape.append(idx[i].size)
elif (isinstance(idx[i], slice)):
shape.append(__slice_len__(idx[i], pu.c2f(A_shape), i))
elif (isinstance(idx[i], arrayfire.Array)):
if idx[i].dtype() is arrayfire.Dtype.b8:
shape.append(int(arrayfire.sum(idx[i])))
else:
shape.append(idx[i].elements())
elif (isinstance(idx[i], arrayfire.index)):
if (idx[i].isspan()):
shape.append(A_shape[i])
else:
af_idx = idx[i].get()
if (af_idx.isBatch):
raise ValueError
if (af_idx.isSeq):
shape.append(arrayfire.seq(af_idx.seq()).size)
else:
shape.append(af_idx.arr_elements())
else:
raise ValueError
return pu.c2f(shape)
