Added a shape test. The op should return the same dynamic shape as the input.

src/include/migraphx/op/prefix_scan_op.hpp

@@ -64,9 +64,9 @@ struct prefix_scan_op : op_name<Derived>
         auto s = inputs.front();
         if(s.dynamic())
         {
-            return {s.type(), s.max_lens()};
-        }
-        if(s.broadcasted())
+            return s;
+        } 
+        else if(s.broadcasted())
         {
             return {s.type(), s.lens()};
         }

test/op_shape_test.cpp

@@ -2084,6 +2084,19 @@ TEST_CASE(prefix_scan_sum)
     }
 }
 
+TEST_CASE(prefix_scan_sum_dyn)
+{
+    {
+        std::vector<migraphx::shape::dynamic_dimension> dd{{5, 8}};
+        migraphx::shape s{migraphx::shape::float_type, dd};
+
+        expect_shape(
+            s,
+            migraphx::make_op("prefix_scan_sum", {{"axis", 0}, {"exclusive", 0}, {"reverse", 0}}),
+            s);
+    }
+}
+
 TEST_CASE(quant_convolution_shape)
 {
     migraphx::shape output{migraphx::shape::int32_type, {4, 4, 1, 1}};

test/ref_ops_test.cpp

@@ -5543,19 +5543,15 @@ TEST_CASE(prefix_scan_sum_dyn_1d)
 {
     migraphx::program p;
     auto* mm = p.get_main_module();
-    std::vector<migraphx::shape::dynamic_dimension> dd{{5, 6}};
+    std::vector<migraphx::shape::dynamic_dimension> dd{{5, 8}};
     migraphx::shape s{migraphx::shape::float_type, dd};
     auto input = mm->add_parameter("X", s);
     mm->add_instruction(migraphx::make_op("prefix_scan_sum", {{"axis", 0}, {"exclusive", false}}), input);
-
-
     p.compile(migraphx::make_target("ref"));
 
     std::vector<float> a = {1, 2, 3, 4, 5, 6};
-    migraphx::parameter_map params0;    
-
-    // auto input = migraphx::literal{s, {1, 2, 3, 4, 5, 6}};
     migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {6}};
+    migraphx::parameter_map params0;    
     params0["X"] = migraphx::argument(input_fixed_shape0, a.data());
 
     auto result = p.eval(params0).back();