namedtensor
namedtensor copied to clipboard
Added proposed multi_index_select version
This PR also fixes the nonzero method default dimension names.
Ouch! I accidentally sent two different functions on the same PR.
I have a proposal that I'd love you can give feedback on it. So it's called multi_index_select(tensor, dims, indices).
Basically, dims refers to the name of dimensions to be indexed by indices; dims must be names of tensor dimensions (no particular order required).
This looks great. the isalphanum stuff could go right in if you wanted to send a separate PR.
One minor thing: I want to move all the asserts to RuntimeError's. That was my fault, but let's do that for future PRs.
A couple requests on multi_index:
-
can you add a longer comment for multi-index-select. It's a bit complicated. In particular I feel like it should take an index_dim argument right? Also unlike other functions the order of dims matters here right?
-
Ideally, this would just work with regular index select. Is there a reason it needs to be a different function? Maybe it just gets called if you pass more than one dim?
- Yes. I need to add documentation there.
I see multi-index-select as a flexible function, where you can say "I have this index, say [[3, 1], [2, 2]], where each element's dimensions correspond to 'dimc' and 'dimb' respectively, and want to index tensor T whose dimensions are (dima, dimb, dimc, dimd)." In this scenario, we would get tensor of dimension (2, dima, dimd).
Ex: tensor.shape -> OrderedDict([("dima", 5), ("dimb", 4), ("dimc", 3), ("dimd", 7)]) indices = [[3, 1], [2, 2]] dims = ('dimc', 'dimb') elements = multi_index_select(tensor, dims, indices) elements.shape -> OrderedDict([("elementsdim", 2), ("dima", 5), ("dimd", 7)])
Is it ways too flexible? I feel like this approach takes great advantage of the abstraction of dimensions proposed by namedtensor.
I could definitively develop other approaches but wanted to propose this idea first.
Yes, I mostly agree with this, I think we are on the same page. However, I would 1) like indices to be a named tensor too not a list, 2) have an argument that says which is the dimension with the indices, 3) not be a different function than index_select.
tensor.shape -> OrderedDict([("dima", 5), ("dimb", 4), ("dimc", 3), ("dimd", 7)]) indices = ntorch.tensor([[3, 1], [2, 2]], names=("elements", "dims") dims = ('dimc', 'dimb') elements = tensor.index_select(dims, indices, index_dim="dims") elements.shape -> OrderedDict([("elements", 2), ("dima", 5), ("dimd", 7)])
where the standard index_select is a special case
tensor.shape -> OrderedDict([("dima", 5), ("dimb", 4), ("dimc", 3), ("dimd", 7)]) indices = ntorch.tensor([[3], [2]], names=("elements") dims = ('dimc') elements = tensor.index_select(dims, indices) elements.shape -> OrderedDict([("elements", 2), ("dima", 5), ("dimd", 7), ("dimb", 4)])
Yes, that makes total sense.
-
Yes, indices is already a
namedtensor
(I failed to say it in the last comment) -
I see. I was missing that! We cannot assume it's the first dimension. For example in this case it would be "elements", which is dim 0:
indices = ntorch.tensor([[3, 1], [2, 2]], names=("elements", "dims"))
However, in this another case it would be "elements", which is dim 1:indices = ntorch.tensor([[3, 1], [2, 2]], names=("dims", "elements"))
-
Right. It doesn't be a different function.
Thanks for the feedback.