Update

Author: ptiede

ext/ReactantStructArraysExt.jl

@@ -53,8 +53,8 @@ function Base.copyto!(
     bc = Broadcast.preprocess(dest, bc)
 
     args = (Reactant.broadcast_to_size(Base.materialize(a), size(bc)) for a in bc.args)
-    @info "HERE"
-    Reactant.TracedUtils.elem_apply_via_while_loop(bc.f, args...; dest=dest)
+    res = Reactant.TracedUtils.elem_apply_via_while_loop(bc.f, args...)
+    copyto!(dest, res)
 
     return dest
 end
@@ -81,11 +81,11 @@ function Base.similar(
     bc::Broadcasted{StructArrayStyle{S,N}}, ::Type{ElType}
 ) where {S<:AbstractReactantArrayStyle,N,ElType}
     bc′ = convert(Broadcasted{S}, bc)
-    if StructArrays.isnonemptystructtype(ElType)
-        StructArrays.buildfromschema(T -> similar(bc′, T), ElType)
-    else
-        similar(bc′, ElType)
-    end
+    alloc(::Type{T}) where {T} = (T <: Complex) ? similar(bc′, T) :
+                                 (StructArrays.isnonemptystructtype(T) ?
+                                  StructArrays.buildfromschema(alloc, T) :
+                                  similar(bc′, T))
+    return alloc(ElType)
 end
 
 Base.@propagate_inbounds function StructArrays._getindex(
@@ -134,11 +134,6 @@ end
     return Base.getfield(obj, field)
 end
 
-function Base.similar(
-    ::Base.Broadcast.Broadcasted{AbstractReactantArrayStyle}, ::Type{Eltype}, dims
-) where {Eltype}
-    return similar(TracedRArray{Eltype}, dims)
-end
 
 # This is to tell StructArrays to leave these array types alone.
 StructArrays.staticschema(::Type{<:Reactant.AnyTracedRArray}) = NamedTuple{()}

src/TracedUtils.jl

@@ -1115,28 +1115,50 @@ function __elem_apply_loop_condition(idx_ref, fn_ref::F, res_ref, args_ref, L_re
     return idx_ref[] < L_ref[]
 end
 
+struct RefFillVector{T}
+    data::T
+end
+
+Base.getindex(rv::RefFillVector, i) = rv.data[]
+Base.broadcastable(x::RefFillVector) = x
+
 function __elem_apply_loop_body(idx_ref, fn_ref::F, res_ref, args_ref, L_ref) where {F}
     args = args_ref[]
     fn = fn_ref[]
     res = res_ref[]
     idx = idx_ref[] + 1
 
-    scalar_args = [@allowscalar(arg[idx]) for arg in args]
+    scalar_args = ntuple(length(args)) do i
+        arg = args[i]
+        return @allowscalar(arg[idx])
+    end
     @allowscalar res[idx] = fn(scalar_args...)
 
     idx_ref[] = idx
     res_ref[] = res
     return nothing
 end
 
-function elem_apply_via_while_loop(f, args::Vararg{Any,Nargs}; dest=nothing) where {Nargs}
-    @assert allequal(size.(args)) "All args must have the same size"
-    L = length(first(args))
+function elem_apply_via_while_loop(f, args::Vararg{Any,Nargs}) where {Nargs}
+    scalar_arg(arg) = arg isa Base.RefValue || !(arg isa AbstractArray)
+    non_ref_args = Tuple(arg for arg in args if !scalar_arg(arg))
+    if !isempty(non_ref_args)
+        @assert allequal(size.(non_ref_args)) "All args must have the same size"
+    end
+    out_size = isempty(non_ref_args) ? () : size(first(non_ref_args))
+    L = isempty(non_ref_args) ? 1 : length(first(non_ref_args))
     # flattening the tensors makes the auto-batching pass work nicer
-    flat_args = [ReactantCore.materialize_traced_array(vec(arg)) for arg in args]
+    flat_args = ntuple(Val(Nargs)) do i
+        arg = args[i]
+        scalar_arg(arg) ? RefFillVector(arg) : ReactantCore.materialize_traced_array(vec(arg))
+    end
 
     # This wont be a mutating function so we can safely execute it once
-    res_tmp = @allowscalar(f([@allowscalar(arg[1]) for arg in flat_args]...))
+    scalar_seed_args = ntuple(Val(Nargs)) do i
+        arg = flat_args[i]
+        @allowscalar(arg[1])
+    end
+    res_tmp = @allowscalar(f(scalar_seed_args...))
 
     # TODO: perhaps instead of this logic, we should have
     # `similar(::TracedRArray, TracedRNumber{T}) where T = similar(::TracedRArray, T)`
@@ -1150,13 +1172,8 @@ function elem_apply_via_while_loop(f, args::Vararg{Any,Nargs}; dest=nothing) whe
     # Before we selected the output container based on the first argument
     # That doesn't work for cases when StructArrays are involved
     # Since this is essentially a broadcast I'm reusing this machinery
-    if isnothing(dest)
-        bc = Base.Broadcast.Broadcasted(f, Tuple(flat_args))
-        result = similar(bc, T_res)
-    else
-        @assert size(dest) == size(first(args)) "dest must have the same size as the input args"
-        result = dest
-    end
+    bc = Base.Broadcast.Broadcasted(f, Tuple(args))
+    result = similar(bc, T_res)
 
     ind_var = Ref(0)
     f_ref = Ref(f)
@@ -1170,7 +1187,7 @@ function elem_apply_via_while_loop(f, args::Vararg{Any,Nargs}; dest=nothing) whe
         (ind_var, f_ref, result_ref, args_ref, limit_ref),
     )
 
-    return ReactantCore.materialize_traced_array(reshape(result, size(first(args))))
+    return ReactantCore.materialize_traced_array(reshape(result, out_size))
 end
 
 function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}