Inconsistent Behaviour in Broadcasting

[willtebbutt]

As of version 0.6.15, the following:

julia> Zygote.gradient(^, 0.0, 0.9)
(0.0, 0.0)

julia> Zygote.gradient((x, y) -> sum(x .^ y), zeros(3), fill(0.9, 3))
([0.0, 0.0, 0.0], [0.0, 0.0, 0.0])

julia> f(x, y) = x ^ y;

julia> Zygote.gradient((x, y) -> sum(f.(x, y)), zeros(3), fill(0.9, 3))
([Inf, Inf, Inf], [NaN, NaN, NaN])

The same code in version 0.6.14:

julia> Zygote.gradient(^, 0.0, 0.9)
(0.0, 0.0)

julia> Zygote.gradient((x, y) -> sum(x .^ y), zeros(3), fill(0.9, 3))
([0.0, 0.0, 0.0], [0.0, 0.0, 0.0])

julia> f(x, y) = x ^ y;

julia> Zygote.gradient((x, y) -> sum(f.(x, y)), zeros(3), fill(0.9, 3))
([0.0, 0.0, 0.0], [0.0, 0.0, 0.0])

the crucial difference being in the last case for each.

If you do this with ForwardDiff, you get

julia> ForwardDiff.gradient(x -> x[1]^x[2], [0.0, 0.9])
2-element Vector{Float64}:
  Inf
 NaN

so this smells to me like ForwardDiff is being invoked for f, whereas individual chain rules are used for sum(x .^ y) because Zygote un-fuses broadcasting where it can (it can't un-fuse f).

@oxinabox @mcabbott @DhairyaLGandhi I'm assuming that this is related to our recent efficiency upgrades?

This is breaking the build in KernelFunctions.jl. We can of course allow the tests to fail for the time being, but it essentially means that we can't use Zygote with one of our kernels.

[willtebbutt]

I guess the culprit is the discrepancy between the version of the rule DiffRules and ChainRules

• https://github.com/JuliaDiff/DiffRules.jl/blob/c90e0bd88e908dc414bf11bb21fc4130800721b0/src/rules.jl#L88
• https://github.com/JuliaDiff/ChainRules.jl/blob/4ad975826ea0639ad709aeb36cc5051b6bf82eaa/src/rulesets/Base/fastmath_able.jl#L170

In the latter we adopt a different convention from the former.

[DhairyaLGandhi]

Hmm, using DiffRules seems to be more "correct". Thoughts?

[willtebbutt]

Unclear to me -- the "subgradient" convention adopted in ChainRules works perfectly in the KernelFunctions use-case. Also, it's what we do for lots of other things (ReLUs, abs, etc). My guess would be it's a useful thing to do?

[willtebbutt]

Hmm, using DiffRules seems to be more "correct". Thoughts?

Hmm, I interpreted this as meaning that you believe the convention adopted in DiffRules is more appropriate than the one in ChainRules. Was I correct in this interpretation?

[oxinabox]

Returning NaN, or Inf for this is just causing pain with little gain, for no real reason.
0.0 is fine

Considering x^y for x=0, and y=0.9
The deriviative wrt x from below is not defined since the function errors if we make a small negative change, so we don't care about that case.
The deriviative wrt x from above is zero.
Thus we might as well say the derivative is zero.

The deriviative wrt y from below is Inf
The deriviative wrt y from above is zero.
I would argue that around if a<=b etc we are largely free to chose which branch we want for a==b (there needs to be some provisos on that statement)
And here chosing zero is the choice that leads to more useful behavour.
Noone want to work with Inf.

This also has the attractive property that if you enumerate all points that the derivative is zero,
i.e. all possible nonboundry local minimal, this gets included.
so if you then check those for the smallest you will find the global minima.
And if you hit this point during gradient desent, you do infact want to stop, so derivative of 0 is useful here also.

Which is where we get the subgradient convention (which I will write up one day, it comes from a discussion I had at a conference.)
We say "If the derivative is not defined, but the subgradient contains zero, then just say the derivative is 0".
Or more directly if taking the derivative from all directions of approach forms a interval that encloses 0, stay the derivative is 0.

[mcabbott]

I think this thread urgently needs a graph:

As x -> 0 from above, the gradient does go to positive real infinity. Whether to return Inf is another question, you could argue that a large finite value from say x = eps() might be nicer.

A few more cases to check, while I was at it. We should be sure to get all of these right.

[antoine-levitt]

How did you obtain these plots? x^y for y non-integer and x negative is impossible to define in a consistent way

[antoine-levitt]

(eg see your 3.14 plot, which does not look like x^3)

[mcabbott]

They allow complex numbers, the dashed lines are the imaginary part when this is nonzero. To get one answer you have to pick a branch, and Julia does. x^(1//3) takes one branch, cbrt(x) takes a different one which is real at negative x.

[antoine-levitt]

It does not pick a branch? (nor should it)

julia> (-2.0)^(1//3)
ERROR: DomainError with -2.0:
Exponentiation yielding a complex result requires a complex argument.
Replace x^y with (x+0im)^y, Complex(x)^y, or similar.

[antoine-levitt]

Ah, you gave it complex numbers, I see. This is different then. In that case the choice is inferred from the fact that zeros are signed:

julia> (-2.0+0.0im)^(1//3)
0.6299605249474366 + 1.0911236359717214im

julia> (-2.0-0.0im)^(1//3)
0.6299605249474366 - 1.0911236359717214im