Ext: add onehotarrays (#1278)

* Ext: add onehotarrays

* fix

* fix prints

* chore: run fmt

* refactor: directly call lu

* Update Project.toml

Author: wsmoses

Project.toml

@@ -36,6 +36,7 @@ LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -53,6 +54,7 @@ ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
 ReactantNNlibExt = "NNlib"
 ReactantOffsetArraysExt = "OffsetArrays"
+ReactantOneHotArraysExt = "OneHotArrays"
 ReactantPythonCallExt = "PythonCall"
 ReactantRandom123Ext = "Random123"
 ReactantSpecialFunctionsExt = "SpecialFunctions"
@@ -80,6 +82,7 @@ LinearAlgebra = "1.10"
 MPI = "0.20"
 NNlib = "0.9.26"
 OffsetArrays = "1"
+OneHotArrays = "0.2.10"
 OrderedCollections = "1"
 PrecompileTools = "1.2"
 Preferences = "1.4"

ext/ReactantOneHotArraysExt.jl

@@ -0,0 +1,23 @@
+module ReactantOneHotArraysExt
+
+using OneHotArrays
+using Reactant
+
+function Reactant.traced_type_inner(
+    @nospecialize(_::Type{OneHotArrays.OneHotArray{T,N,Np1,I}}),
+    seen,
+    @nospecialize(mode::Reactant.TraceMode),
+    @nospecialize(track_numbers::Type),
+    @nospecialize(sharding),
+    @nospecialize(runtime)
+) where {T,N,Np1,I}
+    I2 = Reactant.traced_type_inner(I, seen, mode, track_numbers, sharding, runtime)
+    T2 = if eltype(I2) <: Reactant.TracedRNumber && !(T <: Reactant.TracedRNumber)
+        Reactant.TracedRNumber{T}
+    else
+        T
+    end
+    return OneHotArrays.OneHotArray{T2,N,Np1,I2}
+end
+
+end

test/integration/onehotarrays.jl

@@ -0,0 +1,11 @@
+using Reactant, Test, OneHotArrays, Random
+
+@testset "OneHotArrays" begin
+    m = onehotbatch([10, 20, 30, 10, 10], 10:10:40)
+    r_m = Reactant.to_rarray(m)
+    a = rand(100, 4)
+    r_a = Reactant.to_rarray(a)
+    r_res = @jit r_a * r_m
+    res = a * m
+    @test convert(Array, r_res) ≈ res
+end

test/runtests.jl

@@ -44,6 +44,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "KernelAbstractions" include("integration/kernelabstractions.jl")
         @safetestset "Linear Algebra" include("integration/linear_algebra.jl")
         @safetestset "OffsetArrays" include("integration/offsetarrays.jl")
+        @safetestset "OneHotArrays" include("integration/onehotarrays.jl")
         @safetestset "AbstractFFTs" include("integration/fft.jl")
         @safetestset "SpecialFunctions" include("integration/special_functions.jl")
         @safetestset "Random" include("integration/random.jl")