Non-deterministic `AssertionError` when differentiating `ThreadsX.sum`

[SebastianAment]

Overview

I first posted on Discourse, but have a more Zygote-targeted question now.

I am using the parallel implementation of sum in ThreadsX.jl to accelerate a model of the form

sum(m->m(x), models),

where models is an array of constituent models. This works fine for small tests but gives a non-deterministic AssertionError (see below) on my model. For this reason, I have not been able to identify a MWE but included the stack trace below. Notably, the error occurs in both Zygote's default adjoint and a custom adjoint (see below), on both Intel x86 and Apple ARM processors, and on Julia 1.7.2 and nightly.

I would highly appreciate any help in making this work!

Using Zygote's Default Adjoint

ERROR: LoadError: TaskFailedException
Stacktrace:
  [1] wait
    @ ./task.jl:345 [inlined]
  [2] fetch(t::Task)
    @ Base ./task.jl:360
  [3] (::Zygote.var"#332#334"{Zygote.Context, Task})(#unused#::Nothing)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/base.jl:75
  [4] (::Zygote.var"#2035#back#335"{Zygote.var"#332#334"{Zygote.Context, Task}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [5] macro expansion
    @ ./threadingconstructs.jl:182 [inlined]
  [6] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:158 [inlined]
  [7] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
  [8] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159 [inlined]
  [9] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [10] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:139 [inlined]
 [11] (::typeof(∂(_transduce_assoc_nocomplete)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [12] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:108 [inlined]
 [13] (::typeof(∂(#transduce_assoc#164)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [14] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:94 [inlined]
 [15] (::typeof(∂(transduce_assoc##kw)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [16] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:246 [inlined]
 [17] (::typeof(∂(#foldxt#170)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [18] Pullback
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:246 [inlined]
 [19] (::typeof(∂(foldxt##kw)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [20] Pullback
    @ ~/.julia/packages/ThreadsX/smC2N/src/reduce.jl:19 [inlined]
 [21] (::typeof(∂(#mapreduce#11)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [22] Pullback
    @ ~/.julia/packages/ThreadsX/smC2N/src/reduce.jl:18 [inlined]
 [23] (::typeof(∂(mapreduce##kw)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [24] Pullback
    @ ~/.julia/packages/ThreadsX/smC2N/src/reduce.jl:48 [inlined]
 [25] (::typeof(∂(#sum#14)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [26] Pullback
    @ ~/.julia/packages/ThreadsX/smC2N/src/reduce.jl:48 [inlined]
 [27] (::typeof(∂(sum)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [28] Pullback
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:27 [inlined]
 [29] (::typeof(∂(evaluate)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [30] Pullback
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:37 [inlined]
 [31] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
 [32] (::Zygote.var"#57#58"{typeof(∂(λ))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:41
 [33] gradient(f::Function, args::Library{Vector{Float64}, Vector{Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:76
 [34] model_gradient
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:40 [inlined]
 [35] model_gradient (repeats 2 times)
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:32 [inlined]
 [36] top-level scope
    @ ./timing.jl:241
 [37] include(fname::String)
    @ Base.MainInclude ./client.jl:476
 [38] top-level scope
    @ REPL[2]:1

    nested task error: AssertionError: x === y
    Stacktrace:
      [1] accum
        @ ~/.julia/packages/Zygote/FPUm3/src/lib/lib.jl:36 [inlined]
      [2] accum(::Base.RefValue{Any}, ::Base.RefValue{Any}, ::Nothing, ::Vararg{Any})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/lib.jl:22
      [3] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/phase.jl:59 [inlined]
      [4] (::typeof(∂(evaluate)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
      [5] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/phase.jl:49 [inlined]
      [6] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:27 [inlined]
      [7] (::typeof(∂(λ)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
      [8] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/library.jl:54 [inlined]
      [9] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/core.jl:785 [inlined]
     [10] (::typeof(∂(next)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [11] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/simd.jl:8 [inlined]
     [12] (::typeof(∂(next)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [13] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/core.jl:181 [inlined]
     [14] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:199 [inlined]
     [15] macro expansion
        @ ./simdloop.jl:77 [inlined]
     [16] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/simd.jl:38 [inlined]
     [17] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:198 [inlined]
     [18] (::typeof(∂(_foldl_linear_bulk)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [19] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:192 [inlined]
     [20] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/basics.jl:117 [inlined]
     [21] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:188 [inlined]
     [22] (::typeof(∂(_foldl_array)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [23] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:182 [inlined]
     [24] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:365 [inlined]
     [25] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/threading_utils.jl:56 [inlined]
     [26] (::typeof(∂(_reduce_basecase)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [27] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:150 [inlined]
     [28] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [29] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159 [inlined]
     [30] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [31] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159 [inlined]
     [32] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [33] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159 [inlined]
     [34] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [35] Pullback
        @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159 [inlined]
     [36] (::typeof(∂(_reduce)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [37] Pullback
        @ ./threadingconstructs.jl:178 [inlined]
     [38] (::typeof(∂(λ)))(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface2.jl:0
     [39] (::Zygote.var"#337#338"{Vector{Float64}, typeof(∂(λ))})()
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/base.jl:81
in expression starting at /Users/sebastianament/Documents/SEA/GaussianProcessPhaseMapping/test/gppm.jl:61

Using Custom Adjoint

To accelerate the backward pass, I implemented a custom adjoint for ThreadsX.sum, based on the serial rule here, based on @mcabbott's suggestion.

using ThreadsX
using ChainRules
using ChainRules: RuleConfig, HasReverseMode, rrule_via_ad, ProjectTo, NoTangent, unthunk
# leads to speedup of 0.043804/0.006673 = 6.5 in application with 10 threads on M1 Pro
function ChainRules.rrule(
    config::RuleConfig{>:HasReverseMode}, ::typeof(ThreadsX.sum), f, xs::AbstractArray)
    fx_and_pullbacks = ThreadsX.map(x->rrule_via_ad(config, f, x), xs)
    y = ThreadsX.sum(first, fx_and_pullbacks)

    pullbacks = ThreadsX.map(last, fx_and_pullbacks)

    project = ProjectTo(xs)

    function sum_pullback(ȳ)
        f̄_and_x̄s = ThreadsX.map(f->f(ȳ), pullbacks) # this leads to a non-deterministic error
        # f̄_and_x̄s = map(f->f(ȳ), pullbacks) # NOTE: using the serial map here makes the function execute reliably
        # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
        f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
            NoTangent()
        else
            ThreadsX.sum(first, f̄_and_x̄s)
        end
        x̄s = ThreadsX.map(unthunk ∘ last, f̄_and_x̄s) # project does not support receiving InplaceableThunks
        return NoTangent(), f̄, project(x̄s)
    end
    return y, sum_pullback
end

Using this adjoint, I get the similar non-deterministic error:

ERROR: LoadError: TaskFailedException
Stacktrace:
  [1] wait
    @ ./task.jl:345 [inlined]
  [2] fetch
    @ ./task.jl:360 [inlined]
  [3] _reduce(ctx::Transducers.CancellableDACContext, rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, init::BangBang.NoBang.Empty{Vector}, reducible::Transducers.SizedReducible{SubArray{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}, 1, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Tuple{UnitRange{Int64}}, true}, Int64})
    @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:160
  [4] _reduce(ctx::Transducers.CancellableDACContext, rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, init::BangBang.NoBang.Empty{Vector}, reducible::Transducers.SizedReducible{Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Int64})
    @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159
  [5] _transduce_assoc_nocomplete
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:139 [inlined]
  [6] transduce_assoc(xform::Transducers.Composition{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Map{Type{BangBang.NoBang.SingletonVector}}}, step::Transducers.Completing{typeof(BangBang.append!!)}, init::BangBang.NoBang.Empty{Vector}, coll0::Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}; simd::Val{false}, basesize::Int64, stoppable::Nothing, nestlevel::Nothing)
    @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:108
  [7] foldxt(step::typeof(BangBang.append!!), xform::Transducers.Composition{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Map{Type{BangBang.NoBang.SingletonVector}}}, itr::Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}; init::BangBang.NoBang.Empty{Vector}, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:basesize,), Tuple{Int64}}})
    @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:246
  [8] #tcopy#175
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:354 [inlined]
  [9] #tcollect#181
    @ ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:423 [inlined]
 [10] __map(f::Function, itr::Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ThreadsX.Implementations ~/.julia/packages/ThreadsX/smC2N/src/map.jl:1
 [11] __map
    @ ~/.julia/packages/ThreadsX/smC2N/src/map.jl:1 [inlined]
 [12] _map(::Function, ::Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ThreadsX.Implementations ~/.julia/packages/ThreadsX/smC2N/src/map.jl:13
 [13] _map
    @ ~/.julia/packages/ThreadsX/smC2N/src/map.jl:13 [inlined]
 [14] map(::Function, ::Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ThreadsX.Implementations ~/.julia/packages/ThreadsX/smC2N/src/map.jl:30
 [15] map
    @ ~/.julia/packages/ThreadsX/smC2N/src/map.jl:26 [inlined]
 [16] (::GaussianProcessPhaseMapping.var"#sum_pullback#41"{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, ChainRulesCore.ProjectTo{AbstractArray, NamedTuple{(:elements, :axes), Tuple{Vector{typeof(identity)}, Tuple{Base.OneTo{Int64}}}}}, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}})(ȳ::Vector{Float64})
    @ GaussianProcessPhaseMapping ~/Documents/SEA/GaussianProcessPhaseMapping/src/parallel_sum.jl:16
 [17] ZBack
    @ ~/.julia/packages/Zygote/FPUm3/src/compiler/chainrules.jl:204 [inlined]
 [18] Pullback
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:27 [inlined]
 [19] (::typeof(∂(evaluate)))(Δ::Vector{Float64})
    @ Zygote ./compiler/interface2.jl:0
 [20] Pullback
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:37 [inlined]
 [21] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ./compiler/interface2.jl:0
 [22] (::Zygote.var"#57#58"{typeof(∂(λ))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:41
 [23] gradient(f::Function, args::Library{Vector{Float64}, Vector{Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}})
    @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/interface.jl:76
 [24] model_gradient
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:40 [inlined]
 [25] model_gradient (repeats 2 times)
    @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:32 [inlined]
 [26] top-level scope
    @ ./timing.jl:241
 [27] include(fname::String)
    @ Base.MainInclude ./client.jl:476
 [28] top-level scope
    @ REPL[6]:1

    nested task error: AssertionError: x === y
    Stacktrace:
      [1] accum
        @ ~/.julia/packages/Zygote/FPUm3/src/lib/lib.jl:36 [inlined]
      [2] accum(::Base.RefValue{Any}, ::Base.RefValue{Any}, ::Nothing, ::Vararg{Any})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/lib/lib.jl:22
      [3] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/phase.jl:59 [inlined]
      [4] (::typeof(∂(evaluate)))(Δ::Vector{Float64})
        @ Zygote ./compiler/interface2.jl:0
      [5] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/phase.jl:49 [inlined]
      [6] Pullback
        @ ~/Documents/SEA/GaussianProcessPhaseMapping/src/library.jl:27 [inlined]
      [7] (::typeof(∂(λ)))(Δ::Vector{Float64})
        @ Zygote ./compiler/interface2.jl:0
      [8] (::Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))})(Δ::Vector{Float64})
        @ Zygote ~/.julia/packages/Zygote/FPUm3/src/compiler/chainrules.jl:256
      [9] (::GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}})(f::Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))})
        @ GaussianProcessPhaseMapping ~/Documents/SEA/GaussianProcessPhaseMapping/src/parallel_sum.jl:16
     [10] next(rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, result::BangBang.NoBang.Empty{Vector}, input::Function)
        @ Transducers ~/.julia/packages/Transducers/CnpYX/src/library.jl:54
     [11] next(rf::Transducers.Reduction{Transducers.NoComplete, Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}}, result::BangBang.NoBang.Empty{Vector}, input::Function)
        @ Transducers ~/.julia/packages/Transducers/CnpYX/src/core.jl:785
     [12] macro expansion
        @ ~/.julia/packages/Transducers/CnpYX/src/core.jl:181 [inlined]
     [13] _foldl_array
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:187 [inlined]
     [14] __foldl__
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:182 [inlined]
     [15] foldl_basecase
        @ ~/.julia/packages/Transducers/CnpYX/src/processes.jl:365 [inlined]
     [16] _reduce_basecase(rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, init::BangBang.NoBang.Empty{Vector}, reducible::Transducers.SizedReducible{SubArray{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}, 1, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Tuple{UnitRange{Int64}}, true}, Int64})
        @ Transducers ~/.julia/packages/Transducers/CnpYX/src/threading_utils.jl:56
     [17] _reduce(ctx::Transducers.CancellableDACContext, rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, init::BangBang.NoBang.Empty{Vector}, reducible::Transducers.SizedReducible{SubArray{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}, 1, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Tuple{UnitRange{Int64}}, true}, Int64})
        @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:150
     [18] _reduce(ctx::Transducers.CancellableDACContext, rf::Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, init::BangBang.NoBang.Empty{Vector}, reducible::Transducers.SizedReducible{SubArray{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}, 1, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Tuple{UnitRange{Int64}}, true}, Int64}) (repeats 4 times)
        @ Transducers ~/.julia/packages/Transducers/CnpYX/src/reduce.jl:159
     [19] (::Transducers.var"#165#166"{Transducers.Reduction{Transducers.Map{GaussianProcessPhaseMapping.var"#39#42"{Vector{Float64}}}, Transducers.Reduction{Transducers.Map{Type{BangBang.NoBang.SingletonVector}}, Transducers.BottomRF{Transducers.Completing{typeof(BangBang.append!!)}}}}, BangBang.NoBang.Empty{Vector}, Transducers.CancellableDACContext, Transducers.SizedReducible{SubArray{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}, 1, Vector{Zygote.var"#ad_pullback#47"{Tuple{GaussianProcessPhaseMapping.var"#20#21"{Vector{Float64}}, Phase{Vector{Float64}, Vector{Float64}, Float64, Float64, Float64, GaussianProcessPhaseMapping.GaussianProfile{Float64}}}, typeof(∂(λ))}}, Tuple{UnitRange{Int64}}, true}, Int64}})()
        @ Transducers ./threadingconstructs.jl:178
in expression starting at /Users/sebastianament/Documents/SEA/GaussianProcessPhaseMapping/test/gppm.jl:60

[mcabbott]

The assertion is here:

Zygote.jl/src/lib/lib.jl

Lines 35 to 38 in 4ed3a86

	function accum(x::RefValue, y::RefValue)
	@assert x === y
	return x
	end

Does your problem have mutable structs in it? These sometimes have gradients represented as Ref{Any}, and have bugs, e.g. #1111.

[SebastianAment]

That was indeed the culprit! Curious that I triggered it only though the parallel code. I'll close this issue because of #1111.

Thank you again for your quick help!

[mcabbott]

Would be curious to know better what triggers this. Do you use mutation or just happen to have mutable structs? If you do, is the ordinary evaluation clearly thread-safe?

(Zygote at some point made an attempt to handle mutation, which involved making mutable gradients. It's possible this is less thread-safe than the initial evaluation.)

[SebastianAment]

In this case, mutable was actually unnecessary. It was left over from a previous version of the code, but the evaluation requires no mutation. The only change I made was to remove mutable from the struct and it worked again.

I am open to sharing the code and help hone in on the problem, if you ever have the cycles and desire to do so.

[SebastianAment]

I am not sure how minimal this is, but it is definitely a working example. The code works fine if mutable is removed from the Sin struct. The problematic line is f̄_and_x̄s = ThreadsX.map(f->f(ȳ), pullbacks) in the rrule. Using a serial map in this case also makes the error disappear. I was not able to reproduce the error with Zygote's default adjoint on this example (though the trace above show that it also fails on my model with the same AssertionError).

using ThreadsX
using Zygote
using ChainRules
using ChainRules: RuleConfig, HasReverseMode, rrule_via_ad, ProjectTo, NoTangent, unthunk

function ChainRules.rrule(
    config::RuleConfig{>:HasReverseMode}, ::typeof(ThreadsX.sum), f, xs::AbstractArray)
    fx_and_pullbacks = ThreadsX.map(x->rrule_via_ad(config, f, x), xs)
    y = ThreadsX.sum(first, fx_and_pullbacks)

    pullbacks = ThreadsX.map(last, fx_and_pullbacks)

    project = ProjectTo(xs)

    function sum_pullback(ȳ)
        println("in custom pullback")
        f̄_and_x̄s = ThreadsX.map(f->f(ȳ), pullbacks) # NOTE: this leads to the AssertionError
        # f̄_and_x̄s = map(f->f(ȳ), pullbacks)
        # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
        f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
            NoTangent()
        else
            ThreadsX.sum(first, f̄_and_x̄s)
        end
        x̄s = ThreadsX.map(unthunk ∘ last, f̄_and_x̄s) # project does not support receiving InplaceableThunks
        return NoTangent(), f̄, project(x̄s)
    end
    return y, sum_pullback
end

mutable struct Sin{T}
    α::T
end
function (S::Sin)(x)
    y = zero(x)
    for i in eachindex(S.α)
        y += @. sin(S.α[i] * x)
    end
    return sum(abs2, y)
end
f = Sin(randn(3))

d = 16
x = [randn(d) for _ in 1:10_000]

@time sum(f, x)
@time gradient(f->sum(f, x), f)

@time ThreadsX.sum(f, x)
@time gradient(f->ThreadsX.sum(f, x), f)