fix: tests used older defn

Author: avik-pal

Diff for: src/Tracing.jl

@@ -302,7 +302,7 @@ Base.@nospecializeinfer function traced_type_inner(
         return TracedRArray{T,N}
     elseif mode == TracedSetPath
         if batchmode == BatchNone
-            return T
+            return TracedRArray{T,N}
         elseif batchmode == BatchArray
             if tobatch === nothing
                 TracedRArray{T,N - 1}

Diff for: test/tracing.jl

@@ -56,33 +56,65 @@ using Test
                 (Complex{UInt128}, Complex{UInt128}, TracedRNumber{Complex{UInt128}}),
 
                 # RArray types
-                (ConcreteRArray{Float64,0}, TracedRArray{Float64,0}, TracedRArray{Float64, 0}),
-                (ConcreteRArray{Float64,1}, TracedRArray{Float64,1}, TracedRArray{Float64, 1}),
-                (ConcreteRArray{Float64,2}, TracedRArray{Float64,2}, TracedRArray{Float64, 2}),
-                (ConcreteRArray{Float64,3}, TracedRArray{Float64,3}, TracedRArray{Float64, 3}),
+                (
+                    ConcreteRArray{Float64,0},
+                    TracedRArray{Float64,0},
+                    TracedRArray{Float64,0},
+                ),
+                (
+                    ConcreteRArray{Float64,1},
+                    TracedRArray{Float64,1},
+                    TracedRArray{Float64,1},
+                ),
+                (
+                    ConcreteRArray{Float64,2},
+                    TracedRArray{Float64,2},
+                    TracedRArray{Float64,2},
+                ),
+                (
+                    ConcreteRArray{Float64,3},
+                    TracedRArray{Float64,3},
+                    TracedRArray{Float64,3},
+                ),
 
                 # Array types
-                (Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64}, 1}),
-                (Array{ConcreteRArray{Float64,2},1}, Array{TracedRArray{Float64,2},1}, Array{TracedRArray{Float64,2}, 1}),
+                (Array{Float64,1}, Array{Float64,1}, Array{TracedRNumber{Float64},1}),
+                (
+                    Array{ConcreteRArray{Float64,2},1},
+                    Array{TracedRArray{Float64,2},1},
+                    Array{TracedRArray{Float64,2},1},
+                ),
 
                 # Union types
-                (Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing, TracedRNumber{Int}}),
+                (Union{Nothing,Int}, Union{Nothing,Int}, Union{Nothing,TracedRNumber{Int}}),
                 (
                     Union{Nothing,ConcreteRArray{Float64,1}},
                     Union{Nothing,TracedRArray{Float64,1}},
-                    Union{Nothing, TracedRArray{Float64, 1}}
+                    Union{Nothing,TracedRArray{Float64,1}},
                 ),
 
                 # Ptr types
                 (Ptr{Float64}, Ptr{Float64}, Ptr{TracedRNumber{Float64}}),
-                (Ptr{ConcreteRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}, Ptr{TracedRArray{Float64,1}}),
-                (Core.LLVMPtr{Float64}, Core.LLVMPtr{Float64}, Core.LLVMPtr{TracedRNumber{Float64}}),
+                (
+                    Ptr{ConcreteRArray{Float64,1}},
+                    Ptr{TracedRArray{Float64,1}},
+                    Ptr{TracedRArray{Float64,1}},
+                ),
+                (
+                    Core.LLVMPtr{Float64},
+                    Core.LLVMPtr{Float64},
+                    Core.LLVMPtr{TracedRNumber{Float64}},
+                ),
                 (
                     Core.LLVMPtr{ConcreteRArray{Float64,1}},
                     Core.LLVMPtr{TracedRArray{Float64,1}},
-                    Core.LLVMPtr{TracedRArray{Float64,1}}
+                    Core.LLVMPtr{TracedRArray{Float64,1}},
+                ),
+                (
+                    Base.RefValue{Float64},
+                    Base.RefValue{Float64},
+                    Base.RefValue{TracedRNumber{Float64}},
                 ),
-                (Base.RefValue{Float64}, Base.RefValue{Float64}, Base.RefValue{TracedRNumber{Float64}}),
                 (
                     Base.RefValue{ConcreteRArray{Float64,1}},
                     Base.RefValue{TracedRArray{Float64,1}},
@@ -93,23 +125,30 @@ using Test
                 (Val{0}, Val{0}, Val{0}),
                 (Val{0.5}, Val{0.5}, Val{0.5}),
                 (Val{:x}, Val{:x}, Val{:x}),
-
-
-                (Dict{Int, ConcreteRArray{Float64,0}}, Dict{Int, TracedRArray{Float64,0}}, Dict{Int, TracedRArray{Float64, 0}}),
+                (
+                    Dict{Int,ConcreteRArray{Float64,0}},
+                    Dict{Int,TracedRArray{Float64,0}},
+                    Dict{Int,TracedRArray{Float64,0}},
+                ),
                 (Dict{Int}, Dict{Int}, Dict{Int}),
                 (Dict, Dict, Dict),
-                ((Dict{A, ConcreteRArray{Float64,0}} where A), (Dict{A, TracedRArray{Float64,0}} where A), (Dict{A, TracedRArray{Float64,0}} where A)),
-
-                (Base.Pairs{Symbol, Union{}}, Base.Pairs{Symbol, Union{}}, Base.Pairs{Symbol, Union{}})
+                (
+                    (Dict{A,ConcreteRArray{Float64,0}} where {A}),
+                    (Dict{A,TracedRArray{Float64,0}} where {A}),
+                    (Dict{A,TracedRArray{Float64,0}} where {A}),
+                ),
+                (
+                    Base.Pairs{Symbol,Union{}},
+                    Base.Pairs{Symbol,Union{}},
+                    Base.Pairs{Symbol,Union{}},
+                ),
             ]
                 tracedty = traced_type(
-                    origty, Val(ConcreteToTraced), Union{}
+                    origty, Val(ConcreteToTraced), Union{}, Reactant.BatchNone, nothing
                 )
                 @test tracedty == targetty
 
-                tracedty2 = traced_type(
-                    origty, Val(ConcreteToTraced), ReactantPrimitive
-                )
+                tracedty2 = traced_type(origty, Val(ConcreteToTraced), ReactantPrimitive)
                 @test tracedty2 == targetty
             end
 
@@ -120,21 +159,21 @@ using Test
                 TracedRArray{Float64,3},
             ]
                 @test_throws Union{ErrorException,String} traced_type(
-                      type, Val(ConcreteToTraced), Union{}
+                    type, Val(ConcreteToTraced), Union{}, Reactant.BatchNone, nothing
                 )
             end
         end
         @testset "traced_type exceptions" begin
             @test_throws TracedTypeError Reactant.traced_type(
-                  Real, Val(Reactant.ArrayToConcrete), Union{}
+                Real, Val(Reactant.ArrayToConcrete), Union{}, Reactant.BatchNone, nothing
             )
 
             struct Node
                 x::Vector{Float64}
                 y::Union{Nothing,Node}
             end
             @test_throws NoFieldMatchError traced_type(
-                   Node, Val(ArrayToConcrete), Union{}
+                Node, Val(ArrayToConcrete), Union{}, Reactant.BatchNone, nothing
             )
         end
     end