Symbolic gradient producing surprising results

[Rasmuskh]

I have a question about a surprising behaviour I see when wrapping the symbolic gradient expressions provided by @tullio verbose=1 in their own functions. I was expecting this to give the same gradients as Tullios auto generated functions, but that is not the case in this example. Is this a bug?

using Tullio, Zygote

function convlocal(x, W)
    @tullio c[s, t, c2, b] := x[s+i-1, t+j-1, c1, b] * W[s+i-1, t+j-1, i, j, c1, c2]
    return c
end

function ∇x!(𝛥x, W, 𝛥ℛ)
    @tullio grad=false 𝛥x[(s + i) - 1, (t + j) - 1, c1, b] = 𝛥x[(s + i) - 1, (t + j) - 1, c1, b] + 𝛥ℛ[s, t, c2, b] * conj(W[(s + i) - 1, (t + j) - 1, i, j, c1, c2])
    return 𝛥x
end

function ∇W!(𝛥W, x, 𝛥ℛ)
    @tullio grad=false 𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] = 𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] + 𝛥ℛ[s, t, c2, b] * conj(x[(s + i) - 1, (t + j) - 1, c1, b])
    return 𝛥W
end

kernel_width, kernel_height, ch_in, ch_out = 3, 3, 1, 2
img_width, img_height, batchsize = 10, 10, 30
x = rand(Float32, img_width, img_height, ch_in, batchsize)
W = rand(Float32, img_width, img_height, kernel_width, kernel_height, ch_in, ch_out)
Δy = rand(Float32, img_width-2, img_height-2, ch_out, batchsize)

# method 1: grads computed using ∇x and ∇W
Δx1, ΔW1 = zeros(Float32, size(x)), zeros(Float32, size(W))
∇x!(Δx1, W, Δy)
∇W!(ΔW1, x, Δy)
# method 2: grads computed via Tullios auto generated functions
G = Zygote._pullback(convlocal, x, W)
Δx2 = G[2](Δy)[2]
ΔW2 = G[2](Δy)[3]

@show Δx1 ≈ Δx2
@show ΔW1 ≈ ΔW2

I was expecting Δx1 to match Δx2, but this is not the case (except for when the last dimension of W is singleton...).

Δx1 ≈ Δx2 = false
ΔW1 ≈ ΔW2 = true

[mcabbott]

I think this is a bit like #136, and ∇x! ought to be an error. I haven't thought this one through as carefully, but I do not expect that copying the analytic gradients is going to work here, as index shifts on the LHS aren't meant to work.

I think it ought to be possible to re-write ∇x! so that all index shifts are on the right. But this is a more complicated transformation than the one Tullio knows how to do.

I was concerned it would get the wrong ranges, but from verbose=true this does not seem to be the issue:

julia> G = Zygote._pullback(convlocal, x, W);
┌ Info: left index ranges
│   s = 1:8
│   t = 1:8
│   c2 = Base.OneTo(2)
└   b = Base.OneTo(30)
┌ Info: reduction index ranges
│   i = Base.OneTo(3)
│   j = Base.OneTo(3)
└   c1 = Base.OneTo(1)

julia> ∇x!(Δx1, W, Δy);  # does have the same ranges
┌ Info: left index ranges
│   c1 = Base.OneTo(1)
│   b = Base.OneTo(30)
│   s = Base.OneTo(8)
│   i = 1:3
│   t = Base.OneTo(8)
└   j = 1:3
┌ Info: reduction index ranges
└   c2 = Base.OneTo(2)

The symbolic gradients from verbose=true are:

# @tullio c[s, t, c2, b] := x[s+i-1, t+j-1, c1, b] * W[s+i-1, t+j-1, i, j, c1, c2]
┌ Info: symbolic gradients
│     :(𝛥x[(s + i) - 1, (t + j) - 1, c1, b] = 𝛥x[(s + i) - 1, (t + j) - 1, c1, b] + 𝛥ℛ[s, t, c2, b] * conj(W[(s + i) - 1, (t + j) - 1, i, j, c1, c2]))
└     :(𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] = 𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] + 𝛥ℛ[s, t, c2, b] * conj(x[(s + i) - 1, (t + j) - 1, c1, b]))

And verbose=2 puts these to use like so (cleaned up a little):

┌ Info: ===== Base actor 
function 𝒜𝒸𝓉!(::Type, ℛ::AbstractArray{𝒯}, x, W, 𝒶𝓍s, 𝒶𝓍t, 𝒶𝓍c2, 𝒶𝓍b, 𝒶𝓍i, 𝒶𝓍j, 𝒶𝓍c1, ♻️ = nothing, 💀 = true) where 𝒯
    for b = 𝒶𝓍b, c2 = 𝒶𝓍c2, t = 𝒶𝓍t, s = 𝒶𝓍s
        𝒜𝒸𝒸 = if ♻️ === nothing
                zero(𝒯)
            else
                ℛ[s, t, c2, b]
            end
        for c1 = 𝒶𝓍c1, j = 𝒶𝓍j, i = 𝒶𝓍i
            𝒜𝒸𝒸 = 𝒜𝒸𝒸 + x[(s + i) - 1, (t + j) - 1, c1, b] * W[(s + i) - 1, (t + j) - 1, i, j, c1, c2]
        end
        ℛ[s, t, c2, b] = 𝒜𝒸𝒸
    end
end 

┌ Info: ===== Base actor (symbolic gradient)
function ∇𝒜𝒸𝓉!(::Type, 𝛥x, 𝛥W, 𝛥ℛ::AbstractArray{𝒯}, ℛ, x, W, 𝒶𝓍c1, 𝒶𝓍c2, 𝒶𝓍b, 𝒶𝓍s, 𝒶𝓍i, 𝒶𝓍t, 𝒶𝓍j, ♻️ = nothing, 💀 = true) where 𝒯
    for c1 = 𝒶𝓍c1, j = 𝒶𝓍j, t = 𝒶𝓍t, i = 𝒶𝓍i, s = 𝒶𝓍s, b = 𝒶𝓍b, c2 = 𝒶𝓍c2
        ℰ𝓍1 = conj(W[(s + i) - 1, (t + j) - 1, i, j, c1, c2])
        ℰ𝓍2 = 𝛥ℛ[s, t, c2, b] * ℰ𝓍1
        ℰ𝓍3 = conj(x[(s + i) - 1, (t + j) - 1, c1, b])
        ℰ𝓍4 = 𝛥ℛ[s, t, c2, b] * ℰ𝓍3
        𝛥x[(s + i) - 1, (t + j) - 1, c1, b] = 𝛥x[(s + i) - 1, (t + j) - 1, c1, b] + ℰ𝓍2
        𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] = 𝛥W[(s + i) - 1, (t + j) - 1, i, j, c1, c2] + ℰ𝓍4
    end
end

│   unsafeleft = Symbol[]
│   unsaferight = [:s, :i, :t, :j]

One issue here is that, for 𝒜𝒸𝓉!, it knows that it may break s, t, c2, b among threads. In the gradient, it knows that it is not safe to break s, i, t, j, because writing into 𝛥x[(s + i) - 1, (t + j) - 1, c1, b] may overwrite other values. But in ∇x!, it gets this wrong, and will thread over all LHS indices:

│   leftind = [:c1, :b, :s, :i, :t, :j]

│   unsafeleft = Symbol[]
│   unsaferight = [:s, :i, :t, :j]

However, this isn't the cause of the disagreement here, as it still happens with one thread.

[Rasmuskh]

Thanks for looking into this.
Based on your feedback I wrote a not so elegant function that avoids using shifting indices on the LHS, by doing the summation in a separate step. This produces correct gradients with respect to x.

function ∇x!(W, Δy)
    Δxdummy = device(zeros(Float32, size(x)..., size(W, 6), size(W, 3), size(W, 4)))
    @tullio grad=false Δxdummy[(s + i) - 1, (t + j) - 1, c1, b, c2, i, j] += Δy[s, t, c2, b] * conj(W[(s + i) - 1, (t + j) - 1, i, j, c1, c2])
    Δx = sum(Δxdummy, dims=(5, 6, 7))
    return Δx
end

My motivation for looking into using the symbolic gradient expressions was to speed up local convolutions when working on GPU, but the runtime is about the same for the two approaches, so I might need to work directly with KernelAbstractions.jl for this. By the way here is a link to a related discourse post : link.