`unique()` error in pullback

[DavidKoplow]

I need to calculate the pullback of a function which uses unique(). The variable being differentiated isn't fed into unique() as an argument. On the forward pass everything seems fine, but on the backwards pass Zygote sets the output of unique() to IRTools.Inner.Undefined(). Here is the MWE:

function f1(A)
    @show unique([1])
    A
end
f1(5) #unique returns correct result
Zygote._pullback(f1,5) #unique returns IRTools.Inner.Undefined()

In the above example Zygote is able to correctly calculate both pullbacks, but in the program I am working on, IRTools.Inner.Undefined() leads to an error, which causes the whole thing to crash. For the rest of the code to work, I need to keep unique() in.

Has anyone encountered this issue? What is the fix?

[mcabbott]

The quick fix is:

julia> gradient(x -> sum(abs2, Zygote.forwarddiff(unique, x)), [1,1,2,3,2])
([2, 0, 4, 6, 0],)

julia> gradient(x -> sum(abs2, unique(x)), [1,1,2,3,2])
ERROR: MethodError: no method matching iterate(::IRTools.Inner.Undefined)

Beyond that, someone should write a rule for unique, ideally into ChainRules.jl so that it can be used elsewhere too. While sort has sortperm there is (I think) no obvious built-in function which will tell the indices of the surviving elements in the original array.

[DavidKoplow]

Unfortunately, for this use case I need to use reverse mode AD. Are you sure this is a result of there not being a unique definition in ChainRules.jl? In the MWE the unique call doesn't effect the derivative of A at all.

[mcabbott]

Oh maybe I mis-read the question. You might be looking for Zygote.dropgrad or Zygote.ignore or Zygote.@nograd, to hide some side-computation (for printing, etc) from Zygote completely?

(What I wrote is some mixed-mode thing, it is using ForwardDiff only for one function call, everything outside of that is being handled by Zygote. The same is true of most broadcasting operations, now.)

[DavidKoplow]

Thanks for the help! Zygote.ignore did the trick.

But, could you explain the mixed-mode thing a little more? Let's say I am reverse-differentiating a function A(x), in which B(x) and C(x) are called. Let's say the pullback of C is known but not of B. Are you saying that it is possible to automatically calculate the pullback of B using forwarddiff then use that with the pullback of C to find the gradient of A?

[DhairyaLGandhi]

You might be interested in looking at ChainRules.rrule_via_ad which lets you hook back into ad which is useful for higher order or compositional style functions.