Rules for sortslices, unique
#546
Conversation
|
ReversePropagation and Diffractor failures are unrelated |
| function sortslices_pullback(dy) | ||
| # No actual need to zero this, and if you didn't, then you could widen eltype | ||
| # Also, you could use similar(dy) here not x, same size? | ||
| dx = _zerolike_writeat(x, unthunk(dy), (), inds...) |
There was a problem hiding this comment.
Do we need the unthunk?
If so, should we push it down inside the _zerolike_writeat ?
There was a problem hiding this comment.
I think that ideally, _zerolike_writeat should be upgraded to return an InplaceThunk. And eventually it should be called grad_getindex or something, too.
I'm not sure whether it should handle un-thunking. I guess it wouldn't hurt to add a method. But since most rules at present call unthunk explicitly, maybe it's clearer to call it here?
There was a problem hiding this comment.
The only other place it's called at present is:
https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/Base/array.jl#L364-L373
There was a problem hiding this comment.
Arguably we shouldn't be unthunking if the destination that we are writing into can accept Any.
(but practically that case doesn't really matter since performance is already shot. And likely Zygote will hate that)
There was a problem hiding this comment.
Yes this is far from being a high-performance function!
If you don't take the shortcut above, then not all entries were unique, and thus _zerolike_writeat has to copy dy into dx at some nontrivial indices. So it has to slice up dy, I don't think it can write just one thunk anywhere.
| if dims isa Colon | ||
| xs, ys = vec(x), y | ||
| else | ||
| xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims)) |
There was a problem hiding this comment.
There is an issue open on BlueStyle to remcomment against this
JuliaDiff/BlueStyle#80
If we are going to do this then how do you feel about:
| xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims)) | |
| xs, ys = collect.(eachslice((x, y); dims=dims)) |
There was a problem hiding this comment.
We could avoid this for style. But I think the broadcast is confusing, and perhaps you do too, because I also think it's missing an easy-to-miss dot:
julia> x, y = rand(2,3), rand(2,3);
julia> collect.(eachslice((x, y); dims=1))
ERROR: MethodError: no method matching eachslice(::Tuple{Matrix{Float64}, Matrix{Float64}}; dims=1)
julia> collect.(eachslice.((x, y); dims=1))
(SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}[[0.5119304534786525, 0.6182654562598278, 0.16701230957752622], [0.7959010118362386, 0.9477852004109513, 0.3864
src/rulesets/Base/sort.jl
Outdated
| xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims)) | ||
| end | ||
| mask = isequal.(permutedims(ys), xs) # unique([0.0, -0.0, NaN, NaN]) | ||
| mask .= (mask .== cumsum(mask, dims=1) .== true) |
There was a problem hiding this comment.
is the .== true for handling missing ?
There was a problem hiding this comment.
No, it's that false == 0 satisfies the first ==.
This is a hacky way of writing findfirst(randn(3,3) .> 0; dims=1) as that doesn't exist. I feel like there ought to be a cleverer way like accumulate(xor, mask; dims) or something, but I didn't see it yet.
There was a problem hiding this comment.
Maybe it wants a comment:
| mask .= (mask .== cumsum(mask, dims=1) .== true) | |
| mask .= (mask .== cumsum(mask, dims=1) .== true) # this implements findfirst(mask; dims=1) |
There was a problem hiding this comment.
should we open an issue on JuliaLang/julia and link it here?
There was a problem hiding this comment.
I guess another way to write this is map(findfirst, eachcol(mask)), since we have many slices already.
There was a problem hiding this comment.
Am lazy to wait for another round of CI, so I think I call it good enough for now.
Co-authored-by: Lyndon White <[email protected]>
| end | ||
| mask = isequal.(permutedims(ys), xs) # unique([0.0, -0.0, NaN, NaN]) | ||
| mask .= (mask .== cumsum(mask, dims=1) .== true) | ||
| keep = map(I -> I[1], findall(mask)) |
There was a problem hiding this comment.
| keep = map(I -> I[1], findall(mask)) | |
| keep = map(first, findall(mask)) |
There was a problem hiding this comment.
I don't think this will work:
julia> map(first, findall(randn(3,3) .> 0))
ERROR: iteration is deliberately unsupported for CartesianIndex. Use `I` rather than `I...`, or use `Tuple(I)...`
There was a problem hiding this comment.
I sort-of understand why you can't splat or broadcast it, but it's pretty weird that you can still index it.
| function sortslices_pullback(dy) | ||
| # No actual need to zero this, and if you didn't, then you could widen eltype | ||
| # Also, you could use similar(dy) here not x, same size? | ||
| dx = _zerolike_writeat(x, unthunk(dy), (), inds...) |
There was a problem hiding this comment.
Arguably we shouldn't be unthunking if the destination that we are writing into can accept Any.
(but practically that case doesn't really matter since performance is already shot. And likely Zygote will hate that)
|
The argument against this treatment of The counter is that we don't know about activity. Zygote will want a rule for unique even if you aren't differentiating with respect to its argument, and at present this fails. Sadly, in this case, the calculation done here is quite complicated and would be better skipped. Maybe much more of it should be inside a |
This is interesting as a topic. But I think beyond the scope of this PR. We can change it later. |
Co-authored-by: Lyndon White <[email protected]>
Closes #392