dict value fix (#688)

* dict value fix * Update test/tracing.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
EnzymeAD · Feb 4, 2025 · 21f7660 · 21f7660
1 parent 234d168
commit 21f7660
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/src/Tracing.jl b/src/Tracing.jl
@@ -164,9 +164,19 @@ end
 Base.@nospecializeinfer @inline dict_key(::Type{<:AbstractDict}) = nothing
 Base.@nospecializeinfer @inline dict_key(::Type{<:AbstractDict{K}}) where {K} = K
 Base.@nospecializeinfer @inline dict_value(::Type{<:AbstractDict}) = nothing
-Base.@nospecializeinfer @inline dict_value(
-    ::Type{<:(AbstractDict{K,V} where {K})}
-) where {V} = V
+Base.@nospecializeinfer @inline function dict_value(
+    T::Type{<:(AbstractDict{K,V} where {K})}
+) where {V}
+    if @isdefined(V)
+        V
+    elseif T <: UnionAll
+        dict_value(T.body)
+    elseif T <: Dict && length(T.parameters) >= 2
+        T.parameters[2]
+    else
+        error("Could not get element type of $T")
+    end
+end
 
 Base.@nospecializeinfer function traced_type_inner(
     @nospecialize(T::Type{<:AbstractDict}),

diff --git a/test/tracing.jl b/test/tracing.jl
@@ -138,6 +138,23 @@ end
                     (Dict{A,TracedRArray{Float64,0}} where {A}),
                     (Dict{A,TracedRArray{Float64,0}} where {A}),
                 ),
+                (
+                    (
+                        Dict{
+                            Symbol,NTuple{nsteps,SpectralVariable3D}
+                        } where {nsteps} where {SpectralVariable3D}
+                    ),
+                    (
+                        Dict{
+                            Symbol,NTuple{nsteps,SpectralVariable3D}
+                        } where {nsteps} where {SpectralVariable3D}
+                    ),
+                    (
+                        Dict{
+                            Symbol,NTuple{nsteps,SpectralVariable3D}
+                        } where {nsteps} where {SpectralVariable3D}
+                    ),
+                ),
                 (
                     Base.Pairs{Symbol,Union{}},
                     Base.Pairs{Symbol,Union{}},