From b267d95d04e2ab0f17be8fbc96fec515a1c536f4 Mon Sep 17 00:00:00 2001
From: Frames Catherine White <oxinabox@ucc.asn.au>
Date: Wed, 17 Aug 2022 22:08:50 +0100
Subject: [PATCH] Rewrap Dict Tangents

---
 src/compiler/chainrules.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/compiler/chainrules.jl b/src/compiler/chainrules.jl
index 7c7de8655..1d17ea91d 100644
--- a/src/compiler/chainrules.jl
+++ b/src/compiler/chainrules.jl
@@ -344,4 +344,8 @@ z2d(dx::NamedTuple{L,S}, primal::AbstractDict) where {L,S<:Tuple{Vararg{Union{Nu
   end
 end
 
+function z2d(dx::AbstractDict, primal::T) where T<:Dict
+      return Tangent{T}(Dict(kk => z2d(dvv, primal[kk]) for (kk, dvv) in dx))
+end
+
 z2d(dx::Ref, primal) = z2d(dx[], primal)  # mutable structs