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)