fix: broadcasting with OHA (#1353)

* fix: broadcasting with OHA

* Update ext/ReactantOneHotArraysExt.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix: convert to Array

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.121"
+version = "0.2.122"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/ReactantOneHotArraysExt.jl

@@ -20,4 +20,23 @@ function Reactant.traced_type_inner(
     return OneHotArrays.OneHotArray{T2,N,Np1,I2}
 end
 
+# OneHotArray is a <: AbstractArray{Bool, M} so our usual dispatches don't work
+function Reactant.TracedUtils.broadcast_to_size(
+    r::OneHotArrays.OneHotArray{T,N,Np1,<:Reactant.TracedRArray}, rsize
+) where {T,N,Np1}
+    return Reactant.TracedUtils.broadcast_to_size(
+        Reactant.TracedUtils.materialize_traced_array(r), rsize
+    )
+end
+
+function Reactant.TracedUtils.materialize_traced_array(r::OneHotArrays.OneHotArray)
+    return reshape(r.indices, 1, size(r.indices)...) .== 1:(r.nlabels)
+end
+
+function Base.Array(
+    r::OneHotArrays.OneHotArray{T,N,Np1,<:Reactant.AbstractConcreteArray}
+) where {T,N,Np1}
+    return Array(reshape(Array(r.indices), 1, size(r.indices)...) .== 1:(r.nlabels))
+end
+
 end

src/TracedRArray.jl

@@ -376,7 +376,7 @@ function _setindex_linear!(a::TracedRArray{T,N}, v, indices::AbstractArray) wher
     res = Ops.scatter_setindex(
         a,
         scalar_index_to_cartesian(vec(indices), size(a)),
-        materialize_traced_array(vec(v)),
+        TracedUtils.promote_to(TracedRArray{T,1}, materialize_traced_array(vec(v))),
     )
     set_mlir_data!(a, get_mlir_data(res))
     return a

test/integration/onehotarrays.jl

@@ -9,3 +9,16 @@ using Reactant, Test, OneHotArrays, Random
     res = a * m
     @test convert(Array, r_res) ≈ res
 end
+
+@testset "broadcasting" begin
+    m = onehotbatch([10, 20, 30, 10, 10], 10:10:40)
+    r_m = Reactant.to_rarray(m)
+    x = rand(Float32, 4, 5)
+    x_ra = Reactant.to_rarray(x)
+
+    res_ra = @jit r_m .+ x_ra
+    res = m .+ x
+    @test res_ra ≈ res
+
+    @test Array(r_m) isa Matrix{Bool}
+end