Different jacobian than FiniteDiff for rfft

[judober]

Maybe I'm doing something terribly wrong here. I tried getting the jacobian from an rfft.
I get different results with Zygote compared to a numerical evaluation with FiniteDiff:

using FFTW, FiniteDiff, Zygote

function Compare1()
    # x
    xlen = 100
    xre = zeros(xlen)
    xre[1] = 1

    # Numerical Jacobian
    JacNum = FiniteDiff.finite_difference_jacobian(test1real, xre)

    # Zygote
    JacZyg = jacobian(test1real, xre)[1]

    return (JacNum[1:8,1:8], JacZyg[1:8,1:8]) # compare the first values
end

function test1(xinC)  # irfft
    TAnz = 2 * size(xinC, 1) - 2   # Even
    xoutR = irfft(xinC, TAnz)
end

function test1real(xinR) # Real wrapper
    iseven(length(xinR)) || error("Expect even vector length")
    xinC = xinR[1:2:end] .+ im .* xinR[2:2:end]
    xoutR = test1(xinC)
end

function Compare2()
    # x
    xlen = 100
    xre = zeros(xlen)
    xre[1] = 1

    # Numerical Jacobian
    JacNum = FiniteDiff.finite_difference_jacobian(test2real, xre)

    # Zygote
    JacZyg = jacobian(test2real, xre)[1]

    return (JacNum[1:8,1:8], JacZyg[1:8,1:8]) # compare the first values
end

function test2(xinR) # rfft
    xoutC = rfft(xinR)
end

function test2real(xinR) # Real wrapper
    xoutC = test2(xinR)
    entry(i) = begin
        if iseven(i)
            imag(xoutC[div(i, 2)])
        else
            real(xoutC[div(i, 2, RoundUp)])
        end
    end
    xoutR = [entry(i) for i in 1:2*length(xoutC)]
end

(JacNumIrfft, JacZygIrfft) = Compare1();
(JacNumRfft, JacZygRfft) = Compare2();

JacNumIrfft - JacZygIrfft
JacNumRfft - JacZygRfft

My results:

julia> JacNumIrfft - JacZygIrfft
8×8 Matrix{Float64}:
 2.37582e-11  0.0  0.0102041    0.0          0.0102041    0.0         0.0102041    0.0
 2.37582e-11  0.0  0.0101831   -0.000653778  0.0101203   -0.00130487  0.0100159   -0.0019506
 2.37582e-11  0.0  0.0101203   -0.00130487   0.00987036  -0.00258831  0.00945833  -0.00382926
 2.37582e-11  0.0  0.0100159   -0.0019506    0.00945833  -0.00382926  0.00855192  -0.00556668
 2.37582e-11  0.0  0.00987036  -0.00258831   0.00889101  -0.00500732  0.0073301   -0.0070988
 2.37582e-11  0.0  0.00968424  -0.00321539   0.00817769  -0.00610317  0.00583793  -0.0083691
 2.37582e-11  0.0  0.00945833  -0.00382926   0.0073301   -0.0070988   0.00413044  -0.00933074
 2.37582e-11  0.0  0.00919356  -0.00442739   0.00636214  -0.00797787  0.00227062  -0.00994824

julia> JacNumRfft - JacZygRfft
8×8 Matrix{Float64}:
  0.0   0.0         0.0        0.0        0.0        0.0        0.0        0.0
  0.0   0.0         0.0        0.0        0.0        0.0        0.0        0.0
 -1.0  -0.998027   -0.992115  -0.982287  -0.968583  -0.951057  -0.929776  -0.904827
  0.0   0.0627905   0.125333   0.187381   0.24869    0.309017   0.368125   0.425779
 -1.0  -0.992115   -0.968583  -0.929776  -0.876307  -0.809017  -0.728969  -0.637424
  0.0   0.125333    0.24869    0.368125   0.481754   0.587785   0.684547   0.770513
 -1.0  -0.982287   -0.929776  -0.844328  -0.728969  -0.587785  -0.425779  -0.24869
  0.0   0.187381    0.368125   0.535827   0.684547   0.809017   0.904827   0.968583

To my understanding, they should be approximately the same but aren't.

When comparing the results of rfft:

julia> JacNumRfft
8×8 Matrix{Float64}:
 1.0   1.0         1.0        1.0        1.0        1.0        1.0        1.0
 0.0   0.0         0.0        0.0        0.0        0.0        0.0        0.0
 1.0   0.998027    0.992115   0.982287   0.968583   0.951057   0.929776   0.904827
 0.0  -0.0627905  -0.125333  -0.187381  -0.24869   -0.309017  -0.368125  -0.425779
 1.0   0.992115    0.968583   0.929776   0.876307   0.809017   0.728969   0.637424
 0.0  -0.125333   -0.24869   -0.368125  -0.481754  -0.587785  -0.684547  -0.770513
 1.0   0.982287    0.929776   0.844328   0.728969   0.587785   0.425779   0.24869
 0.0  -0.187381   -0.368125  -0.535827  -0.684547  -0.809017  -0.904827  -0.968583

julia> JacZygRfft
8×8 Matrix{Float64}:
 1.0   1.0        1.0        1.0        1.0        1.0        1.0        1.0
 0.0   0.0        0.0        0.0        0.0        0.0        0.0        0.0
 2.0   1.99605    1.98423    1.96457    1.93717    1.90211    1.85955    1.80965
 0.0  -0.125581  -0.250666  -0.374763  -0.49738   -0.618034  -0.736249  -0.851559
 2.0   1.98423    1.93717    1.85955    1.75261    1.61803    1.45794    1.27485
 0.0  -0.250666  -0.49738   -0.736249  -0.963507  -1.17557   -1.36909   -1.54103
 2.0   1.96457    1.85955    1.68866    1.45794    1.17557    0.851559   0.49738
 0.0  -0.374763  -0.736249  -1.07165   -1.36909   -1.61803   -1.80965   -1.93717

It looks as if there is a factor of two difference between them beginning with row three. The same is true for irfft beginning with the third column.
Indeed:

julia> JacZygRfft ./ [ 1,1, 2,2,2,2,2,2] ≈ JacNumRfft
true

julia> JacZygIrfft .* [ 1,1, 2,2,2,2,2,2]' ≈ JacNumIrfft
true

So maybe an issue with the CainRules definition for AbstractFFTs.jl?

[mcabbott]

No insight but Zygote isn't yet using ChainRules for this, rfft is here:

Zygote.jl/src/lib/array.jl

Lines 848 to 860 in 60f53e7

	# to actually use rfft, one needs to insure that everything
	# that happens in the Fourier domain could've been done in
	# the space domain with real numbers. This means enforcing
	# conjugate symmetry along all transformed dimensions besides
	# the first. Otherwise this is going to result in *very* weird
	# behavior.
	@adjoint function rfft(xs::AbstractArray{<:Real})
	return AbstractFFTs.rfft(xs), function(Δ)
	N = length(Δ)
	originalSize = size(xs,1)
	return (AbstractFFTs.brfft(Δ, originalSize),)
	end
	end

Cc @devmotion who was recently working on JuliaMath/AbstractFFTs.jl#58

[devmotion]

I became convinced (also since otherwise CRTU fails as in the comment above 😄) that the Zygote implementation is incorrect when defining the rules in AbstractFFTs. There an explicit scaling is performed to avoid the problem in this issue: https://github.com/JuliaMath/AbstractFFTs.jl/blob/4f630d6a45b44ff5149eb4ab8ecf60ea0dd510c8/src/chainrules.jl#L25-L32

[trahflow]

I think this can be closed.
(Fixed by correct ChainRules in AbstractFFTs, #1386 and JuliaMath/AbstractFFTs.jl#84)

I verified that the script by @judober above now gives correct results:

julia> JacNumIrfft - JacZygIrfft
8×8 Matrix{Float64}:
 2.37582e-11  0.0   4.75165e-11   0.0          …   4.75165e-11   0.0
 2.37582e-11  0.0  -1.21479e-11   3.54696e-11      5.52809e-11  -2.31852e-11
 2.37582e-11  0.0  -8.25372e-11  -3.87722e-11     -3.04843e-11  -4.88242e-11
 2.37582e-11  0.0   5.52809e-11  -2.31852e-11     -2.45996e-11  -2.25277e-12
 2.37582e-11  0.0   5.77732e-11  -5.3607e-12       8.8986e-11   -2.73702e-11
 2.37582e-11  0.0  -1.65912e-11   2.66676e-11  …  -6.75754e-11  -1.38724e-11
 2.37582e-11  0.0  -3.04843e-11  -4.88242e-11      1.53845e-12   7.5563e-11
 2.37582e-11  0.0  -1.18102e-10   8.44658e-11     -7.21179e-12  -5.09448e-11

julia> JacNumRfft - JacZygRfft
8×8 Matrix{Float64}:
 0.0   0.0           0.0           0.0          …   0.0         0.0
 0.0   0.0           0.0           0.0              0.0         0.0
 0.0   2.01794e-10   6.51309e-9    7.18109e-9       3.8464e-9   5.84926e-9
 0.0   4.16334e-17   5.55112e-17  -5.55112e-17      0.0         5.55112e-17
 0.0   6.51309e-9    5.47064e-9    3.8464e-9       -7.12112e-9  2.40829e-9
 0.0   0.0          -5.55112e-17   0.0          …   0.0         0.0
 0.0   7.18109e-9    3.8464e-9     1.13373e-9       6.38281e-9  2.74298e-9
 0.0  -5.55112e-17   0.0           0.0              0.0         0.0