Use template parameter for mean

src/targets/gpu/jit/reduce.cpp

@@ -164,7 +164,7 @@ struct reduce_compiler : compiler<reduce_compiler>
             auto reduce_elements = get_reduce_elements(ins->inputs());
             auto reduce_type     = ins->inputs().front()->get_shape().type();
             v["reduction"]       = "op::sum{}";
-            std::string mean     = "op::mean{" + std::to_string(reduce_elements) + "}";
+            std::string mean     = "op::mean<" + std::to_string(reduce_elements) + ">{}";
             // Use float accumulator when reduction size is too large for half
             if(reduce_type == shape::half_type and reduce_elements > 16384)
                 v["read"] = "compose(" + mean + ", op::convert_to<float>{})";

src/targets/gpu/kernels/include/migraphx/kernels/ops.hpp

@@ -66,13 +66,22 @@ struct convert_to
     }
 };
 
+template<index_int N>
 struct mean
 {
-    index_int item_num = 1;
     template <class T>
-    MIGRAPHX_DEVICE_CONSTEXPR auto operator()(T x) const
+    MIGRAPHX_DEVICE_CONSTEXPR T operator()(T x) const
     {
-        return x / static_cast<T>(item_num);
+        using type = vec_type<T>;
+        if constexpr(is_floating_point<type>{})
+        {
+            constexpr type d = 1.0 / N;
+            return x * d;
+        }
+        else
+        {
+            return x / static_cast<type>(N);
+        }
     }
 };