Error with Diagonal

[mcabbott]

From JuliaDiff/ChainRules.jl#655:

julia> test_rrule(getindex, Diagonal(rand(3)), 2, :) 
test_rrule: getindex on Diagonal{Float64, Vector{Float64}},Int64,Colon: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:193
  Got exception outside of a @test
  ArgumentError: cannot set off-diagonal entry (2, 1) to a nonzero value (-1.37)
  Stacktrace:
    [1] setindex!
      @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/LinearAlgebra/src/diagonal.jl:129 [inlined]
    [2] setindex!
      @ ./subarray.jl:327 [inlined]
    [3] macro expansion
      @ ./broadcast.jl:974 [inlined]
    [4] macro expansion
      @ ./simdloop.jl:77 [inlined]
    [5] copyto!
      @ ./broadcast.jl:973 [inlined]
    [6] copyto!
      @ ./broadcast.jl:926 [inlined]
    [7] materialize!
      @ ./broadcast.jl:884 [inlined]
    [8] materialize!(dest::SubArray{Float64, 1, Diagonal{Float64, Vector{Float64}}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, false}, bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(+), Tuple{SubArray{Float64, 1, Diagonal{Float64, Vector{Float64}}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, false}, Vector{Float64}}})
      @ Base.Broadcast ./broadcast.jl:881
    [9] ∇getindex!(::Diagonal{Float64, Vector{Float64}}, ::Vector{Float64}, ::Int64, ::Vararg{Any})
      @ ChainRules ~/.julia/dev/ChainRules/src/rulesets/Base/indexing.jl:123
   [10] #1545
      @ ~/.julia/dev/ChainRules/src/rulesets/Base/indexing.jl:73 [inlined]
   [11] add!!(x::Diagonal{Float64, Vector{Float64}}, t::InplaceableThunk{Thunk{ChainRules.var"#1546#1548"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, ChainRules.var"#1545#1547"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}})
      @ ChainRulesCore ~/.julia/packages/ChainRulesCore/ctmSK/src/accumulation.jl:18
   [12] _test_add!!_behaviour(acc::Any, val::Any; kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/check_result.jl:188
   [13] _test_cotangent(accum_cotangent::Any, ad_cotangent::Any, fd_cotangent::Any; check_inferred::Any, kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:300
   [14] (::ChainRulesTestUtils.var"#60#64"{Bool, NamedTuple{(:rtol, :atol), Tuple{Float64, Float64}}})(::Any, ::Vararg{Any})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:226
   [15] #59
      @ ./tuple.jl:612 [inlined]
   [16] BottomRF
      @ ./reduce.jl:81 [inlined]
   [17] _foldl_impl
      @ ./reduce.jl:62 [inlined]
   [18] foldl_impl
      @ ./reduce.jl:48 [inlined]
   [19] mapfoldl_impl(f::typeof(identity), op::Base.var"#59#60"{ChainRulesTestUtils.var"#60#64"{Bool, NamedTuple{(:rtol, :atol), Tuple{Float64, Float64}}}}, nt::Nothing, itr::Base.Iterators.Zip{Tuple{Tuple{NoTangent, Diagonal{Float64, Vector{Float64}}, NoTangent, NoTangent}, Tuple{NoTangent, InplaceableThunk{Thunk{ChainRules.var"#1546#1548"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, ChainRules.var"#1545#1547"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, NoTangent, NoTangent}, Tuple{NoTangent, Diagonal{Float64, Vector{Float64}}, NoTangent, NoTangent}}})
      @ Base ./reduce.jl:44
   [20] mapfoldl(f::Function, op::Function, itr::Base.Iterators.Zip{Tuple{Tuple{NoTangent, Diagonal{Float64, Vector{Float64}}, NoTangent, NoTangent}, Tuple{NoTangent, InplaceableThunk{Thunk{ChainRules.var"#1546#1548"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, ChainRules.var"#1545#1547"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, NoTangent, NoTangent}, Tuple{NoTangent, Diagonal{Float64, Vector{Float64}}, NoTangent, NoTangent}}}; init::Nothing)
      @ Base ./reduce.jl:170
   [21] #foldl#293
      @ ./reduce.jl:193 [inlined]
   [22] foreach(::Function, ::Tuple{NoTangent, Diagonal{Float64, Vector{Float64}}, NoTangent, NoTangent}, ::Tuple{NoTangent, InplaceableThunk{Thunk{ChainRules.var"#1546#1548"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, ChainRules.var"#1545#1547"{Diagonal{Float64, Vector{Float64}}, Vector{Float64}, Tuple{Int64, Colon}}}, NoTangent, NoTangent}, ::Vararg{Tuple})
      @ Base ./tuple.jl:612
   [23] macro expansion
      @ ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:225 [inlined]
   [24] macro expansion
      @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/Test/src/Test.jl:1492 [inlined]
   [25] test_rrule(::RuleConfig, ::Any, ::Any, ::Vararg{Any}; output_tangent::Any, check_thunked_output_tangent::Any, fdm::Any, rrule_f::Any, check_inferred::Bool, fkwargs::NamedTuple, rtol::Real, atol::Real, kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:196
   [26] test_rrule(::RuleConfig, ::Any, ::Any, ::Vararg{Any})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:173
   [27] test_rrule(::Any, ::Vararg{Any}; kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:170
   [28] test_rrule(::Any, ::Any, ::Any, ::Vararg{Any})
      @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/YbVdW/src/testers.jl:169
   [29] top-level scope
      @ REPL[9]:1
   [30] eval
      @ ./boot.jl:370 [inlined]
   [31] eval_user_input(ast::Any, backend::REPL.REPLBackend, mod::Module)
      @ REPL ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:152
   [32] repl_backend_loop(backend::REPL.REPLBackend, get_module::Function)
      @ REPL ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:248
   [33] start_repl_backend(backend::REPL.REPLBackend, consumer::Any; get_module::Function)
      @ REPL ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:233
   [34] run_repl(repl::REPL.AbstractREPL, consumer::Any; backend_on_current_task::Bool)
      @ REPL ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:372
   [35] run_repl(repl::REPL.AbstractREPL, consumer::Any)
      @ REPL ~/.julia/dev/julia/usr/share/julia/stdlib/v1.9/REPL/src/REPL.jl:357
   [36] (::Base.var"#1013#1015"{Bool, Bool, Bool})(REPL::Module)
      @ Base ./client.jl:413
   [37] #invokelatest#2
      @ ./essentials.jl:810 [inlined]
   [38] invokelatest
      @ ./essentials.jl:807 [inlined]
   [39] run_main_repl(interactive::Bool, quiet::Bool, banner::Bool, history_file::Bool, color_set::Bool)
      @ Base ./client.jl:397
   [40] exec_options(opts::Base.JLOptions)
      @ Base ./client.jl:314
   [41] _start()
      @ Base ./client.jl:514
Test Summary:                                                          | Pass  Error  Total  Time
test_rrule: getindex on Diagonal{Float64, Vector{Float64}},Int64,Colon |    5      1      6  2.1s
ERROR: Some tests did not pass: 5 passed, 0 failed, 1 errored, 0 broken.

but it seems to work fine:

julia> y, bk = rrule(getindex, Diagonal(rand(3)), 2, :)
([0.0, 0.5280287668160426, 0.0], ChainRules.var"#getindex_pullback#1544"{Diagonal{Float64, Vector{Float64}}, Tuple{Int64, Colon}}([0.4395686922823513 0.0 0.0; 0.0 0.5280287668160426 0.0; 0.0 0.0 0.8884298121263516], (2, Colon())))

julia> bk([1,2,3])
(NoTangent(), InplaceableThunk(ChainRules.var"#..., Thunk(ChainRules.var"#...)), NoTangent(), NoTangent())

julia> ans[2] |> unthunk
3×3 Diagonal{Float64, Vector{Float64}}:
 0.0   ⋅    ⋅ 
  ⋅   2.0   ⋅ 
  ⋅    ⋅   0.0