Skip to content

Destructuring/iterating a Tangent does not unthunk fields, unlike getindex/getproperty #712

Open

Description

@devmotion

Field access on a Tangent unthunks the field, but iteration does not:

julia> using ChainRulesCore

julia> t = Tangent{Tuple{Float64,Vector{Float64}}}(1.0, @thunk([2.0]));

julia> t[2]          # Vector{Float64}

julia> _, b = t; b   # Thunk

julia> collect(t)    # [1.0, Thunk]

julia> nt = Tangent{NamedTuple{(:a,),Tuple{Vector{Float64}}}}(a=@thunk([1.0]));

julia> nt.a          # Vector{Float64}

julia> (; a) = nt; a # Vector{Float64}

julia> first(nt)     # Thunk

getindex/getproperty call unthunk (src/tangent_types/structural_tangent.jl, since #121), whereas Base.iterate and Base.indexed_iterate forward to the backing without unthunking.

This caused a bug in ChainRules (JuliaDiff/ChainRules.jl#843): the findnz pullbacks did _, _, V̄ = unthunk(Δ). With Zygote, Δ is Tangent{Any}((ZeroTangent(), ZeroTangent(), thunk)), so V̄ was still a thunk and sparse(I, J, V̄, m, n) threw a MethodError.

There are two possible resolutions:

  1. Make iterate/indexed_iterate unthunk, for consistency with getindex/getproperty. This has a downside: positional destructuring would then also compute components the rule discards, e.g. the _ in _, _, V̄ = Δ.
  2. Keep the current behaviour, but document that rules should access the fields they need via getindex/getproperty (or property destructuring (; a) = Δ) and not via positional destructuring or iteration.

Option 2 seems preferable for performance, but the current behaviour is a silent footgun either way. ChainRules.jl#843 now uses indexing (unthunk(Δ)[3]).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions