Fix shape computation

ROCm · Jan 10, 2025 · 925baf1 · 925baf1
1 parent d24f634
commit 925baf1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/targets/gpu/jit/reduce.cpp b/src/targets/gpu/jit/reduce.cpp
@@ -425,8 +425,9 @@ struct fused_reduce_compiler : compiler<fused_reduce_compiler>
         tuning_config tc;
         auto shapes    = to_shapes(ins->inputs());
         tc.problem = to_value(shapes);
+        auto axes           = op.to_value().at("axes").to_vector<std::size_t>();
         auto input_shape = get_input_shape(shapes);
-        auto reduce_shape = get_reduce_shape(shapes);
+        auto reduce_shape = get_reduced_shape(input_shape, axes);
         auto relements = reduce_shape.elements();
         for(auto block_size:{64, 128, 256, 512, 1024})
         {