JuliaDiff / JuliaDiff/ChainRules.jl
Array `getindex` rule unable to handle Zero types and `NotImplemented`
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 475
- Forks
- 98
- PR merge metrics
- No merged PRs in 30d
Description
I've been revisiting https://github.com/FluxML/Zygote.jl/pull/1328 as part of a larger PR, and discovered this behaviour while running https://github.com/FluxML/Zygote.jl/blob/108e5a19d8fa7187f6eaece7a142c48d71dfd0d2/test/chainrules.jl#L275.
MWE:
julia> _, back = rrule(getindex, [1], 1)
(1, ChainRules.var"#getindex_pullback#1601"{Vector{Int64}, Tuple{Int64}, Tuple{NoTangent}}([1], (1,), (NoTangent(),)))
julia> gs = back(@not_implemented("test"))
(NoTangent(), InplaceableThunk(ChainRules.var"#..., Thunk(ChainRules.var"#...)), NoTangent())
julia> unthunk(gs[2])
ERROR: MethodError: Cannot `convert` an object of type Bool to an object of type ChainRulesCore.NotImplemented
Closest candidates are:
convert(::Type{T}, ::T) where T at Base.jl:61
ChainRulesCore.NotImplemented(::Any, ::Any, ::Any) at ~/.julia/packages/ChainRulesCore/a4mIA/src/tangent_types/notimplemented.jl:30
Stacktrace:
[1] fill!(dest::Vector{ChainRulesCore.NotImplemented}, x::Bool)
@ Base ./array.jl:351
[2] _setindex_zero(x::Vector{Int64}, dy::ChainRulesCore.NotImplemented, inds::Int64)
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:104
[3] ∇getindex(x::Vector{Int64}, dy::ChainRulesCore.NotImplemented, inds::Int64)
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:88
[4] (::ChainRules.var"#1603#1605"{Vector{Int64}, ChainRulesCore.NotImplemented, Tuple{Int64}})()
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:73
[5] unthunk
@ ~/.julia/packages/ChainRulesCore/a4mIA/src/tangent_types/thunks.jl:204 [inlined]
[6] unthunk(x::InplaceableThunk{Thunk{ChainRules.var"#1603#1605"{Vector{Int64}, ChainRulesCore.NotImplemented, Tuple{Int64}}}, ChainRules.var"#1602#1604"{Vector{Int64}, ChainRulesCore.NotImplemented, Tuple{Int64}}})
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/a4mIA/src/tangent_types/thunks.jl:237
[7] top-level scope
@ REPL[11]:1
julia> _, back = rrule(getindex, [1], [1])
([1], ChainRules.var"#getindex_pullback#1601"{Vector{Int64}, Tuple{Vector{Int64}}, Tuple{NoTangent}}([1], ([1],), (NoTangent(),)))
julia> gs = back([NoTangent()])
(NoTangent(), InplaceableThunk(ChainRules.var"#..., Thunk(ChainRules.var"#...)), NoTangent())
julia> unthunk(gs[2])
ERROR: MethodError: Cannot `convert` an object of type Bool to an object of type NoTangent
Closest candidates are:
convert(::Type{T}, ::T) where T at Base.jl:61
Stacktrace:
[1] fill!(dest::Vector{NoTangent}, x::Bool)
@ Base ./array.jl:351
[2] _setindex_zero(x::Vector{Int64}, dy::Vector{NoTangent}, inds::Vector{Int64})
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:105
[3] ∇getindex(x::Vector{Int64}, dy::Vector{NoTangent}, inds::Vector{Int64})
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:88
[4] (::ChainRules.var"#1603#1605"{Vector{Int64}, Vector{NoTangent}, Tuple{Vector{Int64}}})()
@ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/indexing.jl:73
[5] unthunk
@ ~/.julia/packages/ChainRulesCore/a4mIA/src/tangent_types/thunks.jl:204 [inlined]
[6] unthunk(x::InplaceableThunk{Thunk{ChainRules.var"#1603#1605"{Vector{Int64}, Vector{NoTangent}, Tuple{Vector{Int64}}}}, ChainRules.var"#1602#1604"{Vector{Int64}, Vector{NoTangent}, Tuple{Vector{Int64}}}})
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/a4mIA/src/tangent_types/thunks.jl:237
[7] top-level scope
@ REPL[14]:1
The lines at fault are https://github.com/JuliaDiff/ChainRules.jl/blob/4ee4ef5ffa86809885a944d8167503b7e2e550fa/src/rulesets/Base/indexing.jl#L104-L105. I would imagine considering x's eltype in the final array would be beneficial, but I'm not familiar enough with all the edge cases to be sure. Maybe the correct solution is to catch this at a higher level.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with test/chainrules.jl around the referenced line and reproduce both MWE cases. Then inspect src/rulesets/Base/indexing.jl at lines 88 and 104-105, focusing on how _setindex_zero handles NotImplemented and NoTangent; done means the pullback can be unthunked for both scalar and array indices without conversion errors.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100