.== and .<= inside Zygote.gradient() are inaccurate on GPU

General Usage
, , , ,
2

What’s going on is this: Zygote computes the gradient of broadcasting using ForwardDiff, at least for GPU arrays (and with some exceptions for functions it knows). That means broadcasting things like Dual.(y) .<= 1:

julia> y .<= 1
3-element MtlVector{Bool, Metal.PrivateStorage}:
 1
 1
 0

julia> using ForwardDiff: Dual

julia> println(Dual.(y, 1) .<= 1)
Bool[1, 0, 0]

(@v1.11) pkg> st ForwardDiff
Status `~/.julia/environments/v1.11/Project.toml`
  [f6369f11] ForwardDiff v1.0.1

The middle 0 (false) is because Dual(1, 1) is regarded as a small perturbation 1.0 + ε, hence above 1. This behaviour is quite new, versions before ForwardDiff 1.0 behaved differently:

julia> println(Dual.(y, 1) .<= 1)
Bool[1, 1, 0]

(jl_c7B1O4) pkg> st ForwardDiff
Status `/private/var/folders/yq/4p2zwd614y59gszh7y9ypyhh0000gn/T/jl_c7B1O4/Project.toml`
⌃ [f6369f11] ForwardDiff v0.10.38

Whether your example is a bug I’m not sure. Zygote has no reason to be differentiating something within println, as this doesn’t contribute to the answer. In fact broadcasting .<= always produces an array of Bool, which isn’t differentiable, which is a second reason for Zygote not to be differentiating that. Nevertheless, Zygote tends to get its fingers into everything. You can tell it not to with @ignore:

julia> let
           f = x -> begin
               Zygote.@ignore println(x .== 1)
               Zygote.@ignore println(x .<= 1)
               return sum(x)
           end

           x = Float32[0,1,2] |> device
           grad = Flux.gradient(f, x)
       end;
Bool[0, 1, 0]
Bool[1, 1, 0]

But whether you should or not depends on what you are doing. If you are optimising some x, you are going to perturb it, and code which depends strongly on the difference between 0.99999994f0 and 1.0000001f0 is then usually a bad idea.

2 Likes