weird issue arising when using exp function with Diagonal matrices

[aarontrowbridge]

Package Version

0.6.49

Julia Version

1.8.1

OS / Environment

ubuntu linux (pop!_os)

Describe the bug

get dot product dimension mismatch error in gradient call when using exp function on Diagonal matrices.

Steps to Reproduce

here is code to reproduce the error:

using Zygote
using LinearAlgebra

d = rand(3) 

D = Diagonal(d)

D.diag == diag(D)
# output: true

function exp_diag_function(D::Diagonal)
    return Diagonal(exp.(diag(D)))
end

function exp_diag_field_name(D::Diagonal)
    return Diagonal(exp.(D.diag))
end

∇_diag_function(x) = Zygote.gradient(x -> tr(exp_diag_function(x * D)), x)

∇_diag_function(1.0)
# output: (1.6395871633504406,)

∇_diag_field_name(x) = Zygote.gradient(x -> tr(exp_diag_field_name(x * D)), x)

∇_diag_field_name(1.0)
# output: ERROR: DimensionMismatch: x and y are of different lengths!

∇_built_in(x) = Zygote.gradient(x -> tr(exp(x * D)), x)

∇_built_in(1.0)
# output: ERROR: DimensionMismatch: x and y are of different lengths!

Expected Results

there should be no error, somehow using the diag field name breaks things vs using the diag method.

Observed Results

dimension mismatch error with native exp function. here is the stack trace:

ERROR: DimensionMismatch: x and y are of different lengths!
Stacktrace:
  [1] dot
    @ ~/packages/julias/julia-1.8/share/julia/stdlib/v1.8/LinearAlgebra/src/generic.jl:866 [inlined]
  [2] dot
    @ ~/packages/julias/julia-1.8/share/julia/stdlib/v1.8/LinearAlgebra/src/generic.jl:856 [inlined]
  [3] #1483
    @ ~/.julia/packages/ChainRules/2ql0h/src/rulesets/Base/arraymath.jl:108 [inlined]
  [4] unthunk
    @ ~/.julia/packages/ChainRulesCore/C73ay/src/tangent_types/thunks.jl:204 [inlined]
  [5] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:105 [inlined]
  [6] map
    @ ./tuple.jl:223 [inlined]
  [7] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:106 [inlined]
  [8] ZBack
    @ ~/.julia/packages/Zygote/dABKa/src/compiler/chainrules.jl:206 [inlined]
  [9] Pullback
    @ ~/projects/Pico.jl/diagonal.jl:36 [inlined]
 [10] (::typeof(∂(#59)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/dABKa/src/compiler/interface2.jl:0
 [11] (::Zygote.var"#60#61"{typeof(∂(#59))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:45
 [12] gradient(f::Function, args::Float64)
    @ Zygote ~/.julia/packages/Zygote/dABKa/src/compiler/interface.jl:97
 [13] ∇_built_in(x::Float64)
    @ Main ~/projects/Pico.jl/diagonal.jl:36
 [14] top-level scope
    @ ~/projects/Pico.jl/diagonal.jl:38

Relevant log output

No response

[mcabbott]

I'm not sure how it reaches that error message, but I think this is what JuliaDiff/ChainRulesCore.jl#446 is meant to fix. Moving some pieces around here's a variant which produces a different error:

julia> Zygote.gradient(x -> tr(Diagonal(exp.(diag(D * x)))), 1.0)
(3.757688003384361,)

julia> Zygote.gradient(x -> tr(Diagonal(exp.((D * x).diag))), 1.0)
ERROR: MethodError: no method matching (::ChainRulesCore.ProjectTo{Diagonal, ...)(::ChainRulesCore.Tangent{Any, NamedTuple{(:diag,), Tuple{Vector{Float64}}}})

julia> using ChainRulesCore
julia> const ArrayOrZero = Union{AbstractArray, AbstractZero};
julia> function (project::ProjectTo{Diagonal})(dx::Tangent{<:Diagonal}) # structural => natural, as in PR
            return dx.diag isa ArrayOrZero ? Diagonal(project.diag(dx.diag)) : dx
        end
julia> function (project::ProjectTo{Diagonal})(dx::Tangent) # same allowing Tangent{Any}
            return dx.diag isa ArrayOrZero ? Diagonal(project.diag(dx.diag)) : dx
        end

julia> Zygote.gradient(x -> tr(Diagonal(exp.((D * x).diag))), 1.0)
(3.757688003384361,)

This structural => natural projection also fixes the error cases above.

[ToucheSir]

The Diagonal constructor invoked by ∇_built_in is more or less equivalent to exp_diag_field_name, so that they behave similarly makes sense. As Michael noted, the issue is that the pullback for D.diag (namely getfield(D, :diag)) creates a structural tangent without an associated type, and subsequent operations don't project that back to a Diagonal.