feat: unbreak NNlib.gather

EnzymeAD · Nov 9, 2024 · 041cdc1 · 041cdc1
1 parent 2208689
commit 041cdc1
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
@@ -698,7 +698,7 @@ end
 
 function broadcast_to_size(arg::T, rsize) where {T<:Number}
     attr = MLIR.IR.DenseElementsAttribute(Base.fill(arg, Tuple(rsize)))
-    return arg = TracedRArray{T,length(rsize)}(
+    return TracedRArray{T,length(rsize)}(
         (), MLIR.IR.result(MLIR.Dialects.stablehlo.constant(; value=attr), 1), rsize
     )
 end
@@ -710,6 +710,11 @@ function broadcast_to_size(arg::TracedRNumber, rsize)
     )
 end
 
+function broadcast_to_size(arg::AnyTracedRArray{T, 0}, rsize) where {T}
+    arg = materialize_traced_array(arg)
+    return broadcast_to_size(TracedRNumber{T}((), arg.mlir_data), rsize)
+end
+
 function broadcast_to_size(arg::AnyTracedRArray, rsize)
     arg = materialize_traced_array(arg)
     size(arg) == rsize && return arg