Incorrect result on function contains adjoint.

[KDr2]

As mentioned at #180 (comment), when a function contains a call to ``adjoint`, ReverseDiff (with latest dependencies) gives a gradient that differs from other AD packages:

This is an MWE:

using ForwardDiff, ReverseDiff
using LinearAlgebra: norm

f(x) = sum(x' .* x)
x = rand(5)

test = ForwardDiff.gradient!(DiffResults.GradientResult(x), f, x)
@show DiffResults.gradient(test)

seedx = rand(eltype(x), size(x))
tp = ReverseDiff.GradientTape(f, seedx)
@show ReverseDiff.gradient!(tp, x)

Run it on the master branch of ReverseDiff, and ForwardDiff=0.10.18, DiffRules=1.0.2, the two results will be the same.
Run it on the master branch of ReverseDiff, and ForwardDiff=0.10.19, DiffRules=1.3.0, the two results will be different.

The problem was found after DiffRules added these two rules lately:

@define_diffrule Base.conj(x)                 = :(  1                                  )
@define_diffrule Base.adjoint(x)              = :(  1                                  )

https://github.com/JuliaDiff/DiffRules.jl/blob/cbf17ea233d16deb6f0a32bca661c63a79126adc/src/rules.jl#L70

[devmotion]

Note that it's not a general problem with adjoint, eg it works fine without broadcasting:

julia> using ForwardDiff, ReverseDiff

julia> f(x) = sum(x' * x);

julia> x = rand(5);

julia> ForwardDiff.gradient(f, x) ≈ 2 * x
true

julia> ReverseDiff.gradient(f, x) ≈ 2 * x
true

whereas

julia> f(x) = sum(x' .* x);

julia> ForwardDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

julia> ReverseDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
false

# but:
julia> ReverseDiff.gradient(f, x) ≈ fill(sum(x), length(x))
true

So I think we should not opt out of the adjoint rules in DiffRules but use these optimizations and only fix the broken cases with broadcasting.

[devmotion]

I guess the first version still works due to this massive specializations for Adjoint:

ReverseDiff.jl/src/derivatives/linalg/arithmetic.jl

Lines 160 to 323 in 01041c8

	mulargvalue(x) = value(x)
	mulargvalue(x::Adjoint) = adjoint(value(adjoint(x)))
	mulargvalue(x::Transpose) = transpose(value(transpose(x)))
	mulargpullvalue!(x) = pull_value!(x)
	mulargpullvalue!(x::Adjoint) = pull_value!(adjoint(x))
	mulargpullvalue!(x::Transpose) = pull_value!(transpose(x))
	# recording pass #
	#----------------#
	@inline function record_mul(x, y, ::Type{D}) where D
	tp = tape(x, y)
	out = track(*(mulargvalue(x), mulargvalue(y)), D, tp)
	cache = (similar(x, D), similar(y, D))
	record!(tp, SpecialInstruction, *, (x, y), out, cache)
	return out
	end
	@inline function record_mul!(out::TrackedArray{V,D}, x, y) where {V,D}
	copyto!(mulargvalue(out), *(mulargvalue(x), mulargvalue(y)))
	cache = (similar(x, D), similar(y, D))
	record!(tape(x, y), SpecialInstruction, *, (x, y), out, cache)
	return out
	end
	for S1 in (:TrackedArray, :TrackedVector, :TrackedMatrix)
	for S2 in (:TrackedArray, :TrackedVector, :TrackedMatrix)
	@eval begin
	LinearAlgebra.:*(x::$(S1){X,D}, y::$(S2){Y,D}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Transpose{<:TrackedReal,<:$(S1){X,D}}, y::Transpose{<:TrackedReal,<:$(S2){Y,D}}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Adjoint{<:TrackedReal,<:$(S1){X,D}}, y::Adjoint{<:TrackedReal,<:$(S2){Y,D}}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Transpose{<:TrackedReal,<:$(S1){X,D}}, y::$(S2){Y,D}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(S1){X,D}, y::Transpose{<:TrackedReal,<:$(S2){Y,D}}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Adjoint{<:TrackedReal,<:$(S1){X,D}}, y::$(S2){Y,D}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(S1){X,D}, y::Adjoint{<:TrackedReal,<:$(S2){Y,D}}) where {X,Y,D} = record_mul(x, y, D)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::$(S1){X,D}, y::$(S2){Y,D}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::Transpose{<:TrackedReal,<:$(S1){X,D}}, y::Transpose{<:TrackedReal,<:$(S2){Y,D}}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::Adjoint{<:TrackedReal,<:$(S1){X,D}}, y::Adjoint{<:Number,<:$(S2){Y,D}}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::Transpose{<:TrackedReal,<:$(S1){X,D}}, y::$(S2){Y,D}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::$(S1){X,D}, y::Transpose{<:TrackedReal,<:$(S2){Y,D}}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::Adjoint{<:TrackedReal,<:$(S1){X,D}}, y::$(S2){Y,D}) where {V,X,Y,D} = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray{V,D}, x::$(S1){X,D}, y::Adjoint{<:TrackedReal,<:$(S2){Y,D}}) where {V,X,Y,D} = record_mul!(out, x, y)
	end
	end
	for T in ARRAY_TYPES
	@eval begin
	LinearAlgebra.:*(x::$(S1){V,D}, y::$(T)) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(T), y::$(S1){V,D}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Transpose{<:Number,<:$(T)}, y::$(S1){V,D}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(S1){V,D}, y::Transpose{<:Number,<:$(T)}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Adjoint{<:Number,<:$(T)}, y::$(S1){V,D}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(S1){V,D}, y::Adjoint{<:Number,<:$(T)}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Transpose{<:Number,<:$(S1){V,D}}, y::$(T)) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(T), y::Transpose{<:Number,<:$(S1){V,D}}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::Adjoint{<:Number,<:$(S1){V,D}}, y::$(T)) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.:*(x::$(T), y::Adjoint{<:Number,<:$(S1){V,D}}) where {V,D} = record_mul(x, y, D)
	LinearAlgebra.mul!(out::TrackedArray, x::$(S1), y::$(T)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::$(T), y::$(S1)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::$(S1), y::Transpose{<:Number,<:$(T)}) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::Transpose{<:Number,<:$(T)}, y::$(S1)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::$(S1), y::Adjoint{<:Number,<:$(T)}) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::Adjoint{<:Number,<:$(T)}, y::$(S1)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::Transpose{<:Number,<:$(S1)}, y::$(T)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::$(T), y::Transpose{<:Number,<:$(S1)}) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::Adjoint{<:Number,<:$(S1)}, y::$(T)) = record_mul!(out, x, y)
	LinearAlgebra.mul!(out::TrackedArray, x::$(T), y::Adjoint{<:Number,<:$(S1)}) = record_mul!(out, x, y)
	end
	end
	end
	# forward pass #
	#--------------#
	@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(*)})
	a, b = instruction.input
	mulargpullvalue!(a)
	mulargpullvalue!(b)
	if instruction.output isa Number
	value!(instruction.output, mulargvalue(a) * mulargvalue(b))
	else
	mul!(mulargvalue(instruction.output), mulargvalue(a), mulargvalue(b))
	end
	return nothing
	end
	# reverse pass #
	#--------------#
	@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(*)})
	a, b = instruction.input
	a_tmp, b_tmp = instruction.cache
	output = instruction.output
	output_deriv = deriv(output)
	reverse_mul!(output, output_deriv, a, b, a_tmp, b_tmp)
	unseed!(output)
	return nothing
	end
	# a * b
	function reverse_mul!(output, output_deriv, a, b, a_tmp, b_tmp)
	istracked(a) && increment_deriv!(a, mul!(a_tmp, output_deriv, transpose(value(b))))
	istracked(b) && increment_deriv!(b, mul!(b_tmp, transpose(value(a)), output_deriv))
	end
	for (f, F) in ((:transpose, :Transpose), (:adjoint, :Adjoint))
	@eval begin
	# a * f(b)
	function reverse_mul!(output, output_deriv, a, b::$F, a_tmp, b_tmp)
	_b = ($f)(b)
	istracked(a) && increment_deriv!(a, mul!(a_tmp, output_deriv, mulargvalue(b)))
	istracked(_b) && increment_deriv!(_b, ($f)(mul!(b_tmp, ($f)(output_deriv), value(a))))
	end
	# f(a) * b
	function reverse_mul!(output, output_deriv, a::$F, b, a_tmp, b_tmp)
	_a = ($f)(a)
	istracked(_a) && increment_deriv!(_a, ($f)(mul!(a_tmp, output_deriv, value(b))))
	istracked(b) && increment_deriv!(b, mul!(b_tmp, ($f)(mulargvalue(a)), ($f)(output_deriv)))
	end
	# f(a) * f(b)
	function reverse_mul!(output, output_deriv, a::$F, b::$F, a_tmp, b_tmp)
	_a = ($f)(a)
	_b = ($f)(b)
	istracked(_a) && increment_deriv!(_a, ($f)(mul!(a_tmp, ($f)(mulargvalue(b)), ($f)(output_deriv))))
	istracked(_b) && increment_deriv!(_b, ($f)(mul!(b_tmp, ($f)(output_deriv), ($f)(mulargvalue(a)))))
	end
	end
	end
	# adjoint(a) * transpose(b)
	function reverse_mul!(output, output_deriv, a::Adjoint, b::Transpose, a_tmp, b_tmp)
	_a = adjoint(a)
	_b = transpose(b)
	if istracked(_a)
	reverse_mul!(output, output_deriv, transpose(_a), b, a_tmp, b_tmp)
	elseif istracked(_b)
	increment_deriv!(_b, transpose(mul!(b_tmp, adjoint(output_deriv), adjoint(mulargvalue(a)))))
	end
	end
	# transpose(a) * adjoint(b)
	function reverse_mul!(output, output_deriv, a::Transpose, b::Adjoint, a_tmp, b_tmp)
	_a = transpose(a)
	_b = adjoint(b)
	if istracked(_b)
	reverse_mul!(output, output_deriv, a, transpose(_b), a_tmp, b_tmp)
	elseif istracked(_a)
	increment_deriv!(_a, transpose(mul!(a_tmp, adjoint(mulargvalue(b)), adjoint(output_deriv))))
	end
	end

😅

I'm sure it's included for a reason but this seems a bit insane and not like a nice solution that generalizes easily to other wrapper types.

[devmotion]

Another observation:

The same problem exists for transpose, and it existed already for DiffRules 1.0.2 (which defined the same rule for transpose):

pkg> add DiffRules@1.0.2, ReverseDiff

julia> using ReverseDiff

julia> x = rand(5);

julia> f(x) = sum(x' .* x);

julia> ReverseDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

julia> ReverseDiff.ForwardDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

julia> f(x) = sum(transpose(x) .* x);

julia> ReverseDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
false

julia> ReverseDiff.ForwardDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

[devmotion]

I think I figured out the main underlying problem here. And a simple fix seems to be:

julia> using ReverseDiff, LinearAlgebra

julia> Base.@propagate_inbounds Base.getindex(v::LinearAlgebra.AdjOrTransAbsVec{<:ReverseDiff.TrackedReal}, i::Int) = v.parent[i-1+first(axes(v.parent)[1])]

julia> f(x) = sum(x' .* x);

julia> x = rand(5);

julia> ReverseDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

julia> f(x) = sum(transpose(x) .* x);

julia> x = rand(5);

julia> ReverseDiff.gradient(f, x) ≈ fill(2 * sum(x), length(x))
true

[devmotion]

A more detailed explanation:

• The Adjoint{<:TrackedReal} and Transpose{<:TrackedReal} correctly hit

  ReverseDiff.jl/src/derivatives/propagation.jl

  Line 35 in 01041c8

  @inline increment_deriv!(t::AbstractArray, x::AbstractArray, i) = increment_deriv!(t[i], x[i])

  (and similarly for decrement_deriv!)
• Indexing into such an array calls https://github.com/JuliaLang/julia/blob/c5f348726cebbe55e169d4d62225c2b1e587f497/stdlib/LinearAlgebra/src/adjtrans.jl#L178 and lines below
• Issue: The definition in LinearAlgebra calls adjoint/transpose (wrapperop(v)) on the retrieved/inserted values but the new rule for adjoint and the old one for transpose lead to a TrackedReal with a derivative of zero!

So it seems to be sufficient to implement https://github.com/JuliaLang/julia/blob/c5f348726cebbe55e169d4d62225c2b1e587f497/stdlib/LinearAlgebra/src/adjtrans.jl#L178-L181 (for completeness probably all getindex/setindex! methods should be implemented) for T<:TrackedReal without wrapperop(v), i.e.

Base.@propagate_inbounds Base.getindex(v::LinearAlgebra.AdjOrTransAbsVec{<:TrackedReal}, i::Int) = v.parent[i-1+first(axes(v.parent)[1])]
Base.@propagate_inbounds Base.getindex(A::LinearAlgebra.AdjOrTransAbsMat{<:TrackedReal}, i::Int, j::Int) = A.parent[j, i]
Base.@propagate_inbounds Base.setindex!(v::LinearAlgebra.AdjOrTransAbsVec{<:TrackedReal}, x, i::Int) = (setindex!(v.parent, x, i-1+first(axes(v.parent)[1])); v)
Base.@propagate_inbounds Base.setindex!(A::LinearAlgebra.AdjOrTransAbsMat{<:TrackedReal}, x, i::Int, j::Int) = (setindex!(A.parent, x, j, i); A)

Edit: See the comment below for a better solution.

[devmotion]

On a second thought, probably it is easier and leads to fewer problems (no inconsistencies, does not become outdated) if we just rely on the defaults in https://github.com/JuliaLang/julia/blob/c5f348726cebbe55e169d4d62225c2b1e587f497/base/number.jl#L211-L213 and the rules for transpose, adjoint and conj are removed from DiffRules. This was already done for identity since it caused problems: JuliaDiff/DiffRules.jl#64

[devmotion]

Fixed by DiffRules 1.3.1.