Rules for `sortslices`, `unique` by mcabbott · Pull Request #546 · JuliaDiff/ChainRules.jl

mcabbott · 2021-11-24T16:45:13Z

Closes #392

oxinabox · 2021-11-30T20:32:58Z

ReversePropagation and Diffractor failures are unrelated

src/rulesets/Base/sort.jl

oxinabox · 2021-11-30T20:35:56Z

src/rulesets/Base/sort.jl

+    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...)


Do we need the unthunk?
If so, should we push it down inside the _zerolike_writeat ?

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?

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

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)

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.

src/rulesets/Base/sort.jl

oxinabox · 2021-11-30T20:41:00Z

src/rulesets/Base/sort.jl

+        if dims isa Colon
+            xs, ys = vec(x), y
+        else
+            xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims))


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:

Suggested change

xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims))

xs, ys = collect.(eachslice((x, y); dims=dims))

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

oxinabox · 2021-11-30T20:50:58Z

src/rulesets/Base/sort.jl

+            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)


is the .== true for handling missing ?

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.

Maybe it wants a comment:

Suggested change

mask .= (mask .== cumsum(mask, dims=1) .== true)

mask .= (mask .== cumsum(mask, dims=1) .== true) # this implements findfirst(mask; dims=1)

should we open an issue on JuliaLang/julia and link it here?

I guess another way to write this is map(findfirst, eachcol(mask)), since we have many slices already.

Am lazy to wait for another round of CI, so I think I call it good enough for now.

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

oxinabox

Basically LGTM

oxinabox · 2021-11-30T20:51:13Z

src/rulesets/Base/sort.jl

+        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))


Suggested change

keep = map(I -> I[1], findall(mask))

keep = map(first, findall(mask))

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)...`

wow, I hate it.

I sort-of understand why you can't splat or broadcast it, but it's pretty weird that you can still index it.

src/rulesets/Base/sort.jl

oxinabox · 2021-11-30T21:22:49Z

src/rulesets/Base/sort.jl

+    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...)


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)

mcabbott · 2021-11-30T21:30:14Z

The argument against this treatment of unique, BTW, is that if the variables are "active", then perhaps they should all be considered unique. As e.g. here: JuliaDiff/ForwardDiff.jl#481 (comment)

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 @thunk... although at present that's ignored anyway.

oxinabox · 2021-11-30T21:54:48Z

The argument against this treatment of unique, BTW, is that if the variables are "active", then perhaps they should all be considered unique. As e.g. here: JuliaDiff/ForwardDiff.jl#481 (comment)

This is interesting as a topic. But I think beyond the scope of this PR. We can change it later.
We should cross check against what Jax does.
Pretty sure they will do this

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

mcabbott added 4 commits November 24, 2021 11:36

unique, take 1

8281539

add shortcut, and tests

912ae5c

sortslices, too

6c95991

fixup

cc9bd43

oxinabox reviewed Nov 30, 2021

View reviewed changes

src/rulesets/Base/sort.jl Outdated Show resolved Hide resolved

oxinabox reviewed Nov 30, 2021

View reviewed changes

src/rulesets/Base/sort.jl Show resolved Hide resolved

oxinabox reviewed Nov 30, 2021

View reviewed changes

Apply 3 suggestions

3f44810

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

oxinabox approved these changes Nov 30, 2021

View reviewed changes

comment

213987e

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

mcabbott merged commit ce78d3d into JuliaDiff:main Nov 30, 2021

mcabbott deleted the unique branch December 1, 2021 10:51

	xs, ys = collect(eachslice(x; dims=dims)), collect(eachslice(y; dims=dims))
	xs, ys = collect.(eachslice((x, y); dims=dims))

	mask .= (mask .== cumsum(mask, dims=1) .== true)
	mask .= (mask .== cumsum(mask, dims=1) .== true) # this implements findfirst(mask; dims=1)

	keep = map(I -> I[1], findall(mask))
	keep = map(first, findall(mask))

Conversation

mcabbott commented Nov 24, 2021

Uh oh!

oxinabox commented Nov 30, 2021

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott Nov 30, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oxinabox left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott commented Nov 30, 2021

Uh oh!

oxinabox commented Nov 30, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mcabbott Nov 30, 2021 •

edited

Loading