Features/914 adv indexing#923
Conversation
…y slice-indexing. UNTESTED
…sition in the index_proxy
|
GPU cluster tests are currently disabled on this Pull Request. |
| """ | ||
| advanced_indexing = False | ||
| if isinstance(key, DNDarray): | ||
| # DNDARRAY CURRENTLY DOES NOT IMPLEMENT the Iterable interface, need to define __iter__() |
There was a problem hiding this comment.
We could implement the Iterable interface for non-distributed DNDarrays, but we need to give some thought to what to fall back to when split is not None.
There was a problem hiding this comment.
Well, iter(...) is just an Iterator over the first dimension, right? So I think
def __iter__(self):
if self.split == 0:
warnings.warn("Iterating over the split dimension is strongly discouraged")
for i in range(len(self)):
yield self[i]would work as expected
| for i, k in enumerate(key): | ||
| if isinstance(k, Iterable) or isinstance(key, DNDarray): | ||
| advanced_indexing = True | ||
| key[i] = factories.array(key[i]) |
There was a problem hiding this comment.
key[i] might be distributed along self.split (or actually newsplit, whatever that is)
There was a problem hiding this comment.
What would newsplit be?
There was a problem hiding this comment.
We don't know yet because we need a sanitized key first. I guess what I'm saying is, we need to be agnostic about k.split at this point. How about: never mind
There was a problem hiding this comment.
Oh no I'm taking this back. tensors and ndarrays have .splitas well
| expand_key[ellipsis_index - len(key) :] = key[ellipsis_index + 1 :] | ||
| key = expand_key | ||
| if add_dims: | ||
| for i, k in reversed(enumerate(key)): |
There was a problem hiding this comment.
TypeError: 'enumerate' object is not reversible...
ClaudiaComito
left a comment
ClaudiaComito
left a comment
There was a problem hiding this comment.
Hi @ben-bou thanks a lot for this! I've reviewed the new private functions, I'll get to the rest later.
| for i, k in reversed(enumerate(key)): | ||
| if k is None: | ||
| key[i] = slice(None) | ||
| arr = arr.expand_dims(i - add_dims + 1) # is the -1 correct? |
There was a problem hiding this comment.
I think __process_key() shouldn't modify arr. torch will take care of creating the extra dimension when we finally index, right?
There was a problem hiding this comment.
You're right, it will. The difference is where the split axis is in the key. key[split] would not be the index for the split axis if there is a newaxis before it. expand_dims is a convenient way to handle this because it's an in-place operation with correct split semantics.
There was a problem hiding this comment.
Another thing to consider is that named tensors are not compatible with newaxis indexing. So if we are to use torch.names to find the output-split, we have to process the newaxis beforehand.
| if k is None: | ||
| key[i] = slice(None) | ||
| arr = arr.expand_dims(i - add_dims + 1) # is the -1 correct? | ||
| add_dims -= 1 | ||
| # expand key to match the number of dimensions of the DNDarray | ||
| key = tuple(key + [slice(None)] * (arr.ndim - len(key))) |
There was a problem hiding this comment.
Really I'm not sure we need to do anything about a None key element at all.
| """ | ||
| Private method for processing keys for indexing. Returns wether advanced indexing is used as well as a processed key and self. | ||
| A processed key: | ||
| - doesn't cotain any ellipses or newaxis |
There was a problem hiding this comment.
this is potentially not correct, see below (i.e. why not leave newaxis in?)
| local_inds = range_proxy[start - offset : stop - offset] # only works if stop - offset > 0 | ||
| local_inds = local_inds[max(offset - start, 0) % step :: step] |
There was a problem hiding this comment.
How about:
local_start = max(start-offset, 0)
local_stop = max(min(stop-offset, range_proxy.stop), 0)
local_inds = range_proxy[local_start:local_stop]
local_inds = local_inds[max(offset - start, 0) % step :: step]
?
There was a problem hiding this comment.
Good point.
So what you're doing by handling the negative values by hand is actually everything indexing the range_proxy would do, right? So there is no reason to sidestep into using a range and we can directly use
local_slice = slice(local_start + max(offset - start, 0) % step, local_stop, step)| range_proxy = range(self.lshape[split]) | ||
| local_inds = range_proxy[start - offset : stop - offset] # only works if stop - offset > 0 | ||
| local_inds = local_inds[max(offset - start, 0) % step :: step] | ||
| if len(local_inds) and stop > offset: |
There was a problem hiding this comment.
If I got local_stop and local_start right, we probably don't need the stop > offset check any more
There was a problem hiding this comment.
we also don't need the local_inds anymore. Checking if the local slice is empty can be done with
if local_stop > 0:
There was a problem hiding this comment.
The condition for empty slice is local_slice.start > local_slice.stop I think.
There was a problem hiding this comment.
One second, is checking for an empty slice needed at all?
In [90]: torch.ones((3, 4, 5))[:, 5:3, :].shape
Out[90]: torch.Size([3, 0, 5])
In [91]: torch.ones((3, 4, 5))[:, 0:0, :].shape
Out[91]: torch.Size([3, 0, 5])This is exactly what the larray of the empty rank should be.
So if that works, is knowing the output_shape beforehand still needed?
At least in the case where key[split] is a slice we can even construct the resulting lshape_map without communication by calculating local_start and local_stop for every offset.
There was a problem hiding this comment.
And if key[split] is a scalar, the split dimension is collapsed and we also know the output-shape and all the lshapes trivially.
Lastly, if advanced indexing is used, the output shape and lshapes are dependent on the key, at least if it's a dndarray.
So if I'm not mistaken, this allows us to get rid of the torch_proxy.
| proxy_key = list(key) | ||
| if advanced_indexing: | ||
| proxy_key = list(key) | ||
| for i, k in reversed(enumerate(key)): |
There was a problem hiding this comment.
reversed(enumerate) returns TypeError
3 participants
Description
Changes proposed:
newaxisindexing and ellipsistorch_proxyBy adding empty dimensions where Advanced indexing should be. In Order to reconstruct the correct shape, Put the shape in the dimensions name.range_proxy.Notes
iter(dndarray)works anywayiter(0-dim dndarray)also works, but shouldn‘tskip ci