fix: broadcasting of closures

Author: avik-pal

Diff for: src/TracedRArray.jl

@@ -246,7 +246,8 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
         invmap[v] = k
     end
 
-    input_shapes = size.(keys(seen_args))
+    keys_seen = [k for k in keys(seen_args) if k isa TracedTypes]
+    input_shapes = size.(keys_seen)
     # by the time we reach here all args must have same size
     @assert allequal(input_shapes) "input shapes are $(input_shapes)"
     OutShape = isempty(seen_args) ? nothing : first(input_shapes)

Diff for: src/utils.jl

@@ -29,7 +29,10 @@ end
 
 function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=false)
     if sizeof(typeof(f)) != 0 || f isa BroadcastFunction
-        return (true, make_mlir_fn(apply, (f, args...), kwargs, name, concretein)[2:end]...)
+        return (
+            true,
+            make_mlir_fn(apply, (f, args...), kwargs, name, concretein; toscalar)[2:end]...,
+        )
     end
 
     N = length(args)

Diff for: test/bcast.jl

@@ -112,3 +112,23 @@ pow(x, n) = x .^ n
 
     @test pow_compiled(x_ra) ≈ pow(x, 2)
 end
+
+struct CustomBCastFunction{X}
+    x::X
+end
+
+(f::CustomBCastFunction)(x::Number) = f.x + x
+
+function custombcast(x)
+    fn = CustomBCastFunction(3.0)
+    return fn.(x)
+end
+
+@testset "Broadcasting closures / functors" begin
+    x = rand(2, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    custombcast_compiled = @compile custombcast(x_ra)
+
+    @test custombcast_compiled(x_ra) ≈ custombcast(x)
+end