Add dynamic support to Shape and Expand

src/include/migraphx/op/multibroadcast.hpp

@@ -29,11 +29,45 @@
 #include <migraphx/dyn_output.hpp>
 #include <migraphx/common.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/errors.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace op {
 
+inline bool is_broadcastable_lens(std::vector<std::size_t> s0, std::vector<std::size_t> s1)
+{
+    if(s0 == s1)
+        return true;
+    if(s0.size() > s1.size())
+        s0.swap(s1);
+    const auto offset = s1.size() - s0.size();
+    return std::equal(s0.begin(), s0.end(), s1.begin() + offset, s1.end(), [](auto a, auto b) {
+        return a == b or a == 1 or b == 1;
+    });
+}
+
+inline std::vector<std::size_t>
+fixed_dyn_dims_lens(const std::vector<shape::dynamic_dimension>& dims)
+{
+    std::vector<std::size_t> target_lens;
+    target_lens.reserve(dims.size());
+    std::transform(dims.begin(), dims.end(), std::back_inserter(target_lens), [](const auto& d) {
+        return shape::static_dim_value(d);
+    });
+    return target_lens;
+}
+
+inline shape broadcast_to_fixed_dims(const shape& s0,
+                                     shape::type_t t,
+                                     const std::vector<shape::dynamic_dimension>& fixed_dims)
+{
+    const auto target_lens = fixed_dyn_dims_lens(fixed_dims);
+    if(is_broadcastable_lens(s0.lens(), target_lens))
+        return make_bcast_shape(s0, compute_broadcasted_lens(s0.lens(), target_lens));
+    return {t, target_lens};
+}
+
 /**
  * Broadcast multiple dimensions between two tensors.
  * Two versions of this operator: 1 input and 2+ inputs.
@@ -71,6 +105,16 @@ struct multibroadcast
 
         if(inputs.size() == 1)
         {
+            if(not output_dyn_dims.empty() and not s0.dynamic())
+            {
+                if(std::all_of(output_dyn_dims.begin(), output_dyn_dims.end(), [](const auto& d) {
+                       return d.is_fixed();
+                   }))
+                {
+                    return broadcast_to_fixed_dims(s0, t, output_dyn_dims);
+                }
+            }
+
             // Symbolic 1-input mode: opt-in via a fully-symbolic output_dyn_dims attribute.
             // Input may be static (bridged via to_symbolic()) or already symbolic.
             // Range-based dynamic input is not allowed.
@@ -128,6 +172,17 @@ struct multibroadcast
             }
             else
             {
+                if(not output_dyn_dims.empty())
+                {
+                    if(std::all_of(output_dyn_dims.begin(),
+                                   output_dyn_dims.end(),
+                                   [](const auto& d) { return d.is_fixed(); }))
+                    {
+                        return broadcast_to_fixed_dims(s0, t, output_dyn_dims);
+                    }
+                    auto bcast_lens = compute_common_lens(inputs);
+                    return make_bcast_shape(s0, bcast_lens);
+                }
                 // output_lens will not be set for 2+ input version
                 auto bcast_lens = compute_common_lens(inputs);
                 return make_bcast_shape(s0, bcast_lens);

src/include/migraphx/shape.hpp

@@ -182,7 +182,8 @@ struct MIGRAPHX_EXPORT shape
             {
                 if(this->sym_expr == other.sym_expr)
                     return *this;
-                return nullopt;
+                if(not other.is_fixed())
+                    return nullopt;
             }
             auto this_interval  = this->get_interval();
             auto other_interval = other.get_interval();

src/onnx/parse_expand.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -50,14 +50,18 @@ struct parse_expand : op_parser<parse_expand>
         else
         {
             const shape& shape_0 = args[0]->get_shape();
+            std::vector<std::size_t> dims;
+            arg_s.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
             if(shape_0.dynamic())
             {
-                MIGRAPHX_THROW(
-                    "PARSE_EXPAND: dynamic input tensor with fixed dims input not supported");
+                shape target_shape{shape_0.type(), dims};
+                auto out_dyn_dims = compute_broadcasted_dyn_dims(shape_0, target_shape);
+                return info.add_instruction(
+                    make_op("multibroadcast", {{"out_dyn_dims", to_value(out_dyn_dims)}}),
+                    args[0],
+                    args[0]);
             }
             const auto& in_lens = shape_0.lens();
-            std::vector<std::size_t> dims;
-            arg_s.visit([&](auto input) { dims.assign(input.begin(), input.end()); });
             auto out_lens = compute_broadcasted_lens(in_lens, dims);
             return info.add_instruction(make_op("multibroadcast", {{"out_lens", out_lens}}),
                                         args[0]);

src/onnx/parse_shape.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -33,8 +33,7 @@ namespace onnx {
 
 /**
  * If static shape input, creates a literal in migraphx.
- * If dynamic shape input, creates a dimensions_of operator in migraphx (runtime evaluation of
- * shape).
+ * If dynamic shape input, creates a compile-time literal with -1 for non-fixed dimensions.
  */
 struct parse_shape : op_parser<parse_shape>
 {
@@ -74,8 +73,16 @@ struct parse_shape : op_parser<parse_shape>
 
         if(input_shape.dynamic())
         {
-            return info.add_instruction(make_op("dimensions_of", {{"start", start}, {"end", end}}),
-                                        args[0]);
+            std::size_t output_ndim = end - start;
+            auto dyn_dims           = input_shape.dyn_dims();
+            std::vector<int64_t> vec_shape(output_ndim);
+            std::transform(
+                dyn_dims.begin() + start,
+                dyn_dims.begin() + end,
+                vec_shape.begin(),
+                [](const auto& dd) { return dd.is_fixed() ? dd.get_interval().max : -1; });
+            return info.add_literal(migraphx::literal{
+                migraphx::shape{migraphx::shape::int64_type, {output_ndim}}, vec_shape});
         }
         else
         {

src/simplify_dyn_ops.cpp

@@ -25,6 +25,7 @@
 #include <migraphx/op/slice.hpp>
 #include <migraphx/op/onehot.hpp>
 #include <migraphx/op/resize.hpp>
+#include <migraphx/op/multibroadcast.hpp>
 #include <migraphx/matcher.hpp>
 #include <migraphx/make_op.hpp>
 #include <migraphx/literal.hpp>
@@ -146,6 +147,19 @@ struct find_static_2in_broadcasts : match::supports_dynamic_shapes
         {
             broadcast_op.from_value({{"out_lens", out_lens}});
         }
+        else if(broadcast_op.name() == "multibroadcast")
+        {
+            const auto& mb = migraphx::any_cast<op::multibroadcast>(broadcast_op);
+            if(not mb.output_dyn_dims.empty())
+            {
+                broadcast_op.from_value(
+                    {{"out_lens", out_lens}, {"out_dyn_dims", to_value(mb.output_dyn_dims)}});
+            }
+            else
+            {
+                broadcast_op.from_value({{"out_lens", out_lens}, {"out_dyn_dims", {}}});
+            }
+        }
         else
         {
             broadcast_op.from_value({{"out_lens", out_lens}, {"out_dyn_dims", {}}});

test/onnx/parse/expand_test.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -69,9 +69,23 @@ TEST_CASE(expand_dyn_input_dyn_output_test)
     EXPECT(p == prog);
 }
 
-TEST_CASE(expand_dyn_input_static_dims_throw)
+TEST_CASE(expand_dyn_input_static_dims_test)
 {
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s(migraphx::shape::float_type, {{3, 8}, {1, 1}, {1, 1}});
+    auto param = mm->add_parameter("x", s);
+    mm->add_literal(
+        migraphx::literal(migraphx::shape{migraphx::shape::int32_type, {3}}, {3, 4, 4}));
+    std::vector<migraphx::shape::dynamic_dimension> out_dyn_dims{{3, 3}, {4, 4}, {4, 4}};
+    auto ret = mm->add_instruction(
+        migraphx::make_op("multibroadcast", {{"out_dyn_dims", migraphx::to_value(out_dyn_dims)}}),
+        param,
+        param);
+    mm->add_return({ret});
+
     migraphx::onnx_options options;
     options.default_dyn_dim_value = {3, 8};
-    EXPECT(test::throws([&] { read_onnx("expand_dyn_input_static_dims_throw.onnx", options); }));
+    auto prog                     = read_onnx("expand_dyn_input_static_dims_throw.onnx", options);
+    EXPECT(p == prog);
 }

test/onnx/parse/shape_dyn_test0.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,9 +29,9 @@ TEST_CASE(shape_dyn_test0)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
+    mm->add_parameter("x", s);
     migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret = mm->add_instruction(migraphx::make_op("dimensions_of", {{"end", 4}}), p0);
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {-1, 4, -1, -1}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;

test/onnx/parse/shape_dyn_test1.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,10 +29,9 @@ TEST_CASE(shape_dyn_test1)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
-    migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret =
-        mm->add_instruction(migraphx::make_op("dimensions_of", {{"start", 2}, {"end", 4}}), p0);
+    mm->add_parameter("x", s);
+    migraphx::shape s_shape{migraphx::shape::int64_type, {2}};
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {-1, -1}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;

test/onnx/parse/shape_dyn_test2.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,10 +29,9 @@ TEST_CASE(shape_dyn_test2)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
-    migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret =
-        mm->add_instruction(migraphx::make_op("dimensions_of", {{"start", 2}, {"end", 4}}), p0);
+    mm->add_parameter("x", s);
+    migraphx::shape s_shape{migraphx::shape::int64_type, {2}};
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {-1, -1}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;

test/onnx/parse/shape_dyn_test3.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,10 +29,9 @@ TEST_CASE(shape_dyn_test3)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
-    migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret =
-        mm->add_instruction(migraphx::make_op("dimensions_of", {{"start", 1}, {"end", 2}}), p0);
+    mm->add_parameter("x", s);
+    migraphx::shape s_shape{migraphx::shape::int64_type, {1}};
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {4}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;

test/onnx/parse/shape_end_oob_test.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,9 +29,9 @@ TEST_CASE(shape_end_oob_test)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
+    mm->add_parameter("x", s);
     migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret = mm->add_instruction(migraphx::make_op("dimensions_of", {{"end", 4}}), p0);
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {-1, 4, -1, -1}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;

test/onnx/parse/shape_start_oob_test.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -29,9 +29,9 @@ TEST_CASE(shape_start_oob_test)
     migraphx::program p;
     auto* mm = p.get_main_module();
     migraphx::shape s{migraphx::shape::float_type, {{1, 4, {1, 4}}, {4, 4}, {2, 4}, {2, 4}}};
-    auto p0 = mm->add_parameter("x", s);
+    mm->add_parameter("x", s);
     migraphx::shape s_shape{migraphx::shape::int64_type, {4}};
-    auto ret = mm->add_instruction(migraphx::make_op("dimensions_of", {{"end", 4}}), p0);
+    auto ret = mm->add_literal(migraphx::literal{s_shape, {-1, 4, -1, -1}});
     mm->add_return({ret});
 
     migraphx::onnx_options options;