fix: materialize LinRange properly (#1389)

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.131"
+version = "0.2.132"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/TracedUtils.jl

@@ -26,6 +26,13 @@ function ReactantCore.materialize_traced_array(x::AbstractRange)
     return Reactant.aos_to_soa(collect(x))
 end
 
+function ReactantCore.materialize_traced_array(r::LinRange)
+    T = Reactant.unwrapped_eltype(r)
+    idxs = Ops.iota(T, [length(r)]; iota_dimension=1)
+    t = idxs ./ r.lendiv
+    return T.((1 .- t) .* r.start .+ t .* r.stop)
+end
+
 function ReactantCore.materialize_traced_array(x::Base.OneTo)
     return Ops.iota(Reactant.unwrapped_eltype(x), [length(x)]; iota_dimension=1)
 end

test/basic.jl

@@ -1347,3 +1347,16 @@ end
     @test @jit(circshift(x_ra, (-3, 2))) ≈ circshift(x, (-3, 2))
     @test @jit(circshift(x_ra, (5, 2))) ≈ circshift(x, (5, 2))
 end
+
+linrange_mat(x1, x2) = Reactant.materialize_traced_array(LinRange(x1, x2, 10024))
+
+@testset "LinRange" begin
+    x1 = 0.0f0
+    x2 = 1.0f0
+    x1_ra = Reactant.to_rarray(x1; track_numbers=Number)
+    x2_ra = Reactant.to_rarray(x2; track_numbers=Number)
+
+    @test @jit(linrange_mat(x1_ra, x2_ra)) ≈ collect(LinRange(x1, x2, 10024))
+    hlo = repr(@code_hlo(linrange_mat(x1_ra, x2_ra)))
+    @test contains(hlo, "stablehlo.iota")
+end