support for N dim arrays, tests

Author: Qfl3x

ext/ReactantFixedSizeArraysExt.jl

@@ -24,9 +24,10 @@ Base.@nospecializeinfer function Reactant.make_tracer(
     mode;
     kwargs...,
 ) where {T,N}
-    return Reactant.make_tracer(
+    shape = size(prev)
+    return reshape(Reactant.make_tracer(
         seen, parent(prev), (path..., 1), mode; kwargs..., track_numbers=Number
-    )
+    ), shape)
 end
 
 end

test/integration/fixedsizearrays.jl

@@ -0,0 +1,17 @@
+
+using Reactant, Test, FixedSizeArrays
+
+fn(x, y) = (2 .* x .- 3) * y'
+
+@testset "FixedSizeArrays" begin
+    @testset "1D" begin
+        x = FixedSizeArray(fill(3.0f0, 100))
+        rx = Reactant.to_rarray(x)
+        @test @jit(fn(rx, rx)) ≈ fn(x, x)
+    end
+    @testset "2D" begin
+        x = FixedSizeArray(fill(3.0f0, (4,5)))
+        rx = Reactant.to_rarray(x)
+        @test @jit(fn(rx, rx)) ≈ fn(x, x)
+    end
+end

test/memory.jl

@@ -0,0 +1,11 @@
+using Reactant, Test
+
+fn(x,y) = sin.(x) .+ cos.(y)
+
+@testset "Memory test" begin
+    x = Memory{Float32}(fill(2.0f0, 10))
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(fn(x_ra,x_ra)) ≈ fn(x,x)
+end
+

test/runtests.jl

@@ -37,6 +37,9 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Cluster Detection" include("cluster_detector.jl")
         @safetestset "Config" include("config.jl")
         @safetestset "Batching" include("batching.jl")
+        if isdefined(Base, :Memory)
+            @safetestset "Memory" include("memory.jl")
+        end
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"
@@ -51,6 +54,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Python" include("integration/python.jl")
         @safetestset "Optimisers" include("integration/optimisers.jl")
         @safetestset "FillArrays" include("integration/fillarrays.jl")
+        @safetestset "FixedSizeArrays" include("integration/fixedsizearrays.jl")
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "neural_networks"