fix: broadcasted type casting (#156)

* fix: broadcasted type casting

* fix: optimisers compilation

* chore: apply formatting

* fix: handle case where the results get optimized out

* test: add tests for type casting

* fix: type restrict conversion to TypeCast

Author: avik-pal

src/Compiler.jl

@@ -29,9 +29,13 @@ function create_result(tocopy::T, path, result_stores) where {T}
 end
 
 function create_result(tocopy::ConcreteRArray{T,N}, path, result_stores) where {T,N}
-    restore = result_stores[path]
-    delete!(result_stores, path)
-    return :(ConcreteRArray{$T,$N}($restore, $(tocopy.shape)))
+    if haskey(result_stores, path)
+        restore = result_stores[path]
+        delete!(result_stores, path)
+        return :(ConcreteRArray{$T,$N}($restore, $(tocopy.shape)))
+    end
+    # We will set the data for this later
+    return :(ConcreteRArray{$T,$N}($(tocopy.data), $(tocopy.shape)))
 end
 
 function create_result(tocopy::Array{T,N}, path, result_stores) where {T,N}
@@ -67,7 +71,9 @@ function create_result(tocopy::D, path, result_stores) where {K,V,D<:AbstractDic
 end
 
 function create_result(
-    tocopy::Union{Int,AbstractFloat,AbstractString,Nothing,Type,Symbol}, path, result_stores
+    tocopy::Union{Integer,AbstractFloat,AbstractString,Nothing,Type,Symbol},
+    path,
+    result_stores,
 )
     return Meta.quot(tocopy)
 end

src/TracedRArray.jl

@@ -338,6 +338,18 @@ for (jlop, hloop) in (
     end
 end
 
+struct TypeCast{T<:Number} <: Function end
+
+function (::TypeCast{T})(x::TracedRArray{T2,0}) where {T,T2}
+    return promote_to(TracedRArray{T,0}, x)
+end
+
+elem_apply(::Type{T}, x::TracedRArray{T}) where {T<:Number} = x
+function elem_apply(::Type{T}, x::TracedRArray{T2}) where {T<:Number,T2<:Number}
+    # Special Path to prevent going down a despecialized path
+    return elem_apply(TypeCast{T}(), x)
+end
+
 function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
     all(iszero ∘ ndims, args) && return f(args...)
 

test/compile.jl

@@ -27,4 +27,20 @@ Base.sum(x::NamedTuple{(:a,),Tuple{T}}) where {T<:Reactant.TracedRArray} = (; a=
 
         @test fworld(a, b) ≈ ones(2, 2) * 10
     end
+
+    @testset "type casting & optimized out returns" begin
+        a = Reactant.ConcreteRArray(rand(2, 10))
+
+        ftype1(x) = Float64.(x)
+        ftype2(x) = Float32.(x)
+
+        ftype1_compiled = @compile ftype1(a)
+        ftype2_compiled = @compile ftype2(a)
+
+        @test ftype1_compiled(a) isa Reactant.ConcreteRArray{Float64,2}
+        @test ftype2_compiled(a) isa Reactant.ConcreteRArray{Float32,2}
+
+        @test ftype1_compiled(a) ≈ Float64.(a)
+        @test ftype2_compiled(a) ≈ Float32.(a)
+    end
 end