WIP FixedSizeArrays support

Author: Qfl3x

Project.toml

@@ -47,6 +47,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
+FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 
 [sources]
 ReactantCore = {path = "lib/ReactantCore"}
@@ -57,6 +58,7 @@ ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
 ReactantDLFP8TypesExt = "DLFP8Types"
 ReactantFillArraysExt = "FillArrays"
+ReactantFixedSizeArraysExt = "FixedSizeArrays"
 ReactantFloat8sExt = "Float8s"
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
@@ -82,6 +84,7 @@ EnumX = "1"
 Enzyme = "0.13.74"
 EnzymeCore = "0.8.13"
 FillArrays = "1.13"
+FixedSizeArrays = "1.2.0"
 Float8s = "0.1"
 Functors = "0.5"
 GPUArraysCore = "0.2"

ext/ReactantFixedSizeArraysExt.jl

@@ -0,0 +1,34 @@
+module ReactantFixedSizeArraysExt
+
+using FixedSizeArrays
+using Reactant
+using Reactant: TracedRArray, TracedRNumber, Ops
+using ReactantCore: ReactantCore
+
+function Reactant.traced_type_inner(
+    @nospecialize(_::Type{FixedSizeArrays.FixedSizeArray{T, N, Memory{I}}}),
+    seen,
+    @nospecialize(mode::Reactant.TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T, N, I}
+    T2 = Reactant.TracedRNumber{T}
+    I2 = Reactant.TracedRNumber{I}
+    return FixedSizeArrays.FixedSizeArray{T2, N, Memory{I2}}
+end
+
+Base.@nospecializeinfer function Reactant.make_tracer(
+    seen,
+    @nospecialize(prev::FixedSizeArrays.FixedSizeArray{T, N, Memory{I}}),
+    @nospecialize(path),
+    mode; kwargs...
+) where {T, N, I}
+    return FixedSizeArrays.FixedSizeArray(
+        Reactant.make_tracer(
+            seen, parent(prev), (path..., 1), mode; kwargs..., track_numbers=Number
+        )
+    )
+end
+    
+end

src/Tracing.jl

@@ -1148,6 +1148,10 @@ function make_tracer(
     @nospecialize(runtime = nothing),
     kwargs...,
 )
+    @show prev
+    @show path
+    @show seen
+    @show typeof(prev)
     return make_tracer_unknown(
         seen, prev, path, mode; track_numbers, sharding, runtime, kwargs...
     )

test/Project.toml

@@ -9,6 +9,7 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -28,6 +29,7 @@ Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"