Update src/Ops.jl

EnzymeAD · Mar 8, 2025 · b14a1e1 · b14a1e1
1 parent 25a1124
commit b14a1e1
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/Ops.jl b/src/Ops.jl
@@ -2366,14 +2366,18 @@ Applies a reduction function `fn` along the specified `dimensions` of input `x`,
 
     result_type = mlir_type(TracedRArray{T,length(reduced_shape)}, reduced_shape)
 
-    sample_inputs = [Reactant.ConcretePJRTNumber(T(0)), Reactant.ConcretePJRTNumber(T(0))]
+    sample_inputs = [
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0),
+        Reactant.TracedUtils.promote_to(TracedRNumber{T}, 0)
+    ]
 
     func =
         Reactant.TracedUtils.make_mlir_fn(
             fn,
             (sample_inputs),
             (),
-            "reduce_fn";
+            "reduce_fn",
+            false;
             args_in_result=:none,
             return_dialect=:stablehlo,
         ).f