From 5f028a46f58808525784fdb36c3c0db8bf4b2da9 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 13 Mar 2022 19:25:10 +0100 Subject: [PATCH] Fix problems related to unthunk_tangent for IdDict Co-authored-by: Brian Chen --- src/compiler/chainrules.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/compiler/chainrules.jl b/src/compiler/chainrules.jl index 08428c190..e851893f4 100644 --- a/src/compiler/chainrules.jl +++ b/src/compiler/chainrules.jl @@ -1,6 +1,10 @@ @inline unthunk_tangent(x::AbstractThunk) = wrap_chainrules_output(unthunk(x)) @inline unthunk_tangent(x::AbstractArray{<:AbstractThunk}) = map(unthunk_tangent, x) unthunk_tangent(d::IdDict) = IdDict([unthunk_tangent(k) => unthunk_tangent(v) for (k, v) in d]) +function ChainRulesCore.rrule(::typeof(unthunk_tangent), d::IdDict) + unthunk_iddict_pullback(_) = (NoTangent(), ChainRulesCore.@not_implemented "unthunking IdDict") + return d, unthunk_iddict_pullback +end @non_differentiable unthunk_tangent(::IdDict)