-
-
Notifications
You must be signed in to change notification settings - Fork 5.7k
Improve inferability of slicedim
#20154
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
base/abstractarraymath.jl
Outdated
|
|
||
| _slice_inds(new_inds::Tuple, inds::Tuple, d, i) = | ||
| _slice_inds((new_inds..., length(new_inds) + 1 == d ? i : first(inds)), tail(inds), d, i) | ||
| _slice_inds(new_inds::Tuple, inds::Tuple{}, d, i) = new_inds |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could also be
setindex(x::Tuple, n, v) = _setindex((), x, n, v)
_setindex(y::Tuple, x::Tuple, n, v) = _setindex((y..., length(y) + 1 == n ? v : first(x)), tail(x), n, v)
_setindex(y::Tuple, x::Tuple{}, n, v) = yand live in tuple.jl. Or do we have something like this somewhere already?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Seems like a good idea. You'd want the v to come before n in the argument order, though, for consistency with setindex!.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe briefer would be ntuple(j->ifelse(j==d,i,inds[d]), Val{N}), but overall I like the idea of having an internal tuple-setindex, it makes it more readable.
I'd also point out that you may be able to get full inferrability of the output in most cases by manually allocating the output array and then using copy!. As long as the indices-tuple is homogenous, the indices-tuple of the output array differs by being of length N-1 rather than N. This is the same type as tail would produce, so you could define
drop_dth{N,I}(inds::NTuple{N,I}, d) = _drop_dth(tail(inds), (), d, inds...)
_drop_dth(ref, out, d, ind, inds...) = length(out)==d+1 ? _drop_dth(ref, out, d, inds...) : _drop_dth(ref, (out..., ind), d, inds...)
_drop_dth{M,I}(::NTuple{M,I}, out, d)::NTuple{M,I} = out # M = N-1drop_dth is type-unstable at intermediate stages, but the output tuple has known type. Pass out as the indices-tuple to similar, and voila, the return array is inferrable.
(Might need a fallback for non-homogeneous indices tuples, but AFAIK there are no such arrays in existence.)
base/abstractarraymath.jl
Outdated
| # index A[:,:,...,i,:,:,...] where "i" is in dimension "d" | ||
|
|
||
| _slice_inds(new_inds::Tuple, inds::Tuple, d, i) = | ||
| _slice_inds((new_inds..., length(new_inds) + 1 == d ? i : first(inds)), tail(inds), d, i) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd considerevaluate ifelse here since ? will give you one branch per dimension. Might also consider @inlineing this.
base/abstractarraymath.jl
Outdated
|
|
||
| _slice_inds(new_inds::Tuple, inds::Tuple, d, i) = | ||
| _slice_inds((new_inds..., length(new_inds) + 1 == d ? i : first(inds)), tail(inds), d, i) | ||
| _slice_inds(new_inds::Tuple, inds::Tuple{}, d, i) = new_inds |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Seems like a good idea. You'd want the v to come before n in the argument order, though, for consistency with setindex!.
We support |
|
Seems almost worth having a separate code path for this case, e.g., require an extra argument or something. But that's not pretty, so maybe it's better to live with the type-instability. |
0ea52db to
e894075
Compare
|
Introduced @timholy Does |
e894075 to
e1423c5
Compare
|
Hmm, seems likely. This version looks fine to me. |
|
Should |
base/tuple.jl
Outdated
|
|
||
| # returns new tuple; N.B.: becomes no-op if i is out-of-bounds | ||
| setindex(x::Tuple, v, i::Integer) = _setindex((), x, v, i::Integer) | ||
| _setindex(y::Tuple, r::Tuple, v, i::Integer) = |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is the one to inline.
e1423c5 to
e36e14e
Compare
|
Added |
slicedimslicedim
|
Oh, there is also a |
|
Interesting, the The specialized code also did not allow a "too large" dimension as long as indexing with 1, so could be made type-stable. But for consistency, probably it should either always be allowed or always disallowed. |
|
Travis failure on Mac looks unrelated. |
Good news:Doing some more benchmarking I am now pretty convinced that the generic code is significantly faster than the specialized code for Bad news:If we want to preserve the current behavior, I could add back a specialization for |
|
Added a specialization so that this PR only does what the title says and does not change any return types. |
| B[l] = A[j + index_offset] | ||
| l += 1 | ||
| end | ||
| # preserve some special behavior |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
was this tested very carefully before?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No, it's not tested at all. And I didn't feel like adding a test now because the desired behavior might actually be redefined shortly depending on how we resolve #20233.
|
Bump – this seems like a strict improvement. |
|
Barring any objections, I'm going to merge this tomorrow. |
An attempt to fix #20141. Not perfect, but an improvement, and should do no harm.
E.g. with this PR:
On master, they are all just
::Any.