JuliaDiff / JuliaDiff/ChainRules.jl

Array `getindex` rule unable to handle Zero types and `NotImplemented`

Open
#697 7 comments 0 reactions 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.