From 9f0ca43bc4c0c1be5247eb1220c3d9203af60555 Mon Sep 17 00:00:00 2001
From: Uros Petkovic <urpetkov@amd.com>
Date: Wed, 7 Oct 2026 14:00:39 +0200
Subject: [PATCH 01/10] Fixing multibroadcast issue

---
 CHANGELOG.md                    |  1 +
 src/simplify_reshapes.cpp       |  9 ++++++---
 test/simplify_reshapes_test.cpp | 24 ++++++++++++++++++++++++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a446e6bfaa6..0622ba42add 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ Full documentation for MIGraphX is available at
 
 ### Resolved issues
 
+* Fixed `simplify_reshapes` rewriting an element-expanding shape-transform chain through a shared reduction, which could corrupt the reduction's other consumers and produce invalid broadcast shapes.
 * Fixed non-standard GPU output layouts being exposed to integrations that require packed tensors (#5345).
 * Fixed mixed `uint8` and `int8` quantization (#5075).
 * Fixed flash-decoding rebuild for fused `@literal`/`@outline` and extra score-shaped `@param` inputs (#5114).
diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index d2eeecb6e0a..16af504a46d 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -282,9 +282,12 @@ struct find_op_shape_transform_op
                return not contains({"multibroadcast", "broadcast", "contiguous"}, x->name());
            }))
             return false;
-        // A chain that changes the element count is only valid for a reduction
-        return is_reduce(x_ins) or
-               input_ins->get_shape().elements() == x_ins->get_shape().elements();
+        if(input_ins->get_shape().elements() == x_ins->get_shape().elements())
+            return true;
+        // An element-expanding chain can be moved through a reduction only when
+        // the reduction is not shared. The rewrite replaces x_ins globally, so
+        // changing its shape would invalidate its other consumers.
+        return is_reduce(x_ins) and x_ins->outputs().size() == 1;
     }
 
     static bool matches_op(instruction_ref ins)
diff --git a/test/simplify_reshapes_test.cpp b/test/simplify_reshapes_test.cpp
index d2a9002ddbe..204ff8d72f5 100644
--- a/test/simplify_reshapes_test.cpp
+++ b/test/simplify_reshapes_test.cpp
@@ -6494,4 +6494,28 @@ TEST_CASE(op_shape_transform_shadowed_broadcast)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(op_shape_transform_shared_reduce_expanding_chain)
+{
+    migraphx::module m1;
+    {
+        auto image =
+            m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale =
+            m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
+        auto mean =
+            m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
+        auto unsqueeze =
+            m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean);
+        auto expanded = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}), unsqueeze);
+        auto mul = m1.add_instruction(migraphx::make_op("mul"), expanded, scale);
+        auto shared = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), mean);
+        m1.add_return({mul, shared});
+    }
+    auto m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

From e5b5773cb23d0473b007b424581e5300edb764ef Mon Sep 17 00:00:00 2001
From: Uros Petkovic <urpetkov@amd.com>
Date: Wed, 7 Oct 2026 14:05:22 +0200
Subject: [PATCH 02/10] Fixing clang format

---
 test/simplify_reshapes_test.cpp | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/test/simplify_reshapes_test.cpp b/test/simplify_reshapes_test.cpp
index 204ff8d72f5..6d9f59d2081 100644
--- a/test/simplify_reshapes_test.cpp
+++ b/test/simplify_reshapes_test.cpp
@@ -6498,17 +6498,13 @@ TEST_CASE(op_shape_transform_shared_reduce_expanding_chain)
 {
     migraphx::module m1;
     {
-        auto image =
-            m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
-        auto scale =
-            m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
-        auto mean =
-            m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
-        auto unsqueeze =
-            m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean);
-        auto expanded = m1.add_instruction(
+        auto image = m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale = m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
+        auto mean = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
+        auto unsqueeze = m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean);
+        auto expanded  = m1.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}), unsqueeze);
-        auto mul = m1.add_instruction(migraphx::make_op("mul"), expanded, scale);
+        auto mul    = m1.add_instruction(migraphx::make_op("mul"), expanded, scale);
         auto shared = m1.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), mean);
         m1.add_return({mul, shared});

From 936a5cb17094cf9c8039d688d36b56b8feea73d1 Mon Sep 17 00:00:00 2001
From: Uros Petkovic <urpetkov@amd.com>
Date: Wed, 7 Oct 2026 14:15:35 +0200
Subject: [PATCH 03/10] Adding PR number

---
 CHANGELOG.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0622ba42add..9d6cebcffcf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ Full documentation for MIGraphX is available at
 
 ### Resolved issues
 
-* Fixed `simplify_reshapes` rewriting an element-expanding shape-transform chain through a shared reduction, which could corrupt the reduction's other consumers and produce invalid broadcast shapes.
+* Fixed `simplify_reshapes` rewriting an element-expanding shape-transform chain through a shared reduction, which could corrupt the reduction's other consumers and produce invalid broadcast shapes (#5360).
 * Fixed non-standard GPU output layouts being exposed to integrations that require packed tensors (#5345).
 * Fixed mixed `uint8` and `int8` quantization (#5075).
 * Fixed flash-decoding rebuild for fused `@literal`/`@outline` and extra score-shaped `@param` inputs (#5114).

From 21ab3436a76d80d0a7f864448b87bdfdd4627a3e Mon Sep 17 00:00:00 2001
From: Uros Petkovic <urpetkov@amd.com>
Date: Wed, 7 Oct 2026 17:30:46 +0200
Subject: [PATCH 04/10] Adding more changes

---
 CHANGELOG.md                    |  2 +-
 src/simplify_reshapes.cpp       | 17 +++++----
 test/simplify_reshapes_test.cpp | 61 +++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9d6cebcffcf..5dc21d6d74a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,7 @@ Full documentation for MIGraphX is available at
 
 ### Resolved issues
 
-* Fixed `simplify_reshapes` rewriting an element-expanding shape-transform chain through a shared reduction, which could corrupt the reduction's other consumers and produce invalid broadcast shapes (#5360).
+* Fixed `simplify_reshapes` generating invalid inverse shape transforms when moving an element-expanding shape-transform chain through a shared reduction (#5360).
 * Fixed non-standard GPU output layouts being exposed to integrations that require packed tensors (#5345).
 * Fixed mixed `uint8` and `int8` quantization (#5075).
 * Fixed flash-decoding rebuild for fused `@literal`/`@outline` and extra score-shaped `@param` inputs (#5114).
diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index 16af504a46d..42bf7fc9e06 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -284,9 +284,7 @@ struct find_op_shape_transform_op
             return false;
         if(input_ins->get_shape().elements() == x_ins->get_shape().elements())
             return true;
-        // An element-expanding chain can be moved through a reduction only when
-        // the reduction is not shared. The rewrite replaces x_ins globally, so
-        // changing its shape would invalidate its other consumers.
+        // Expanding chains cannot restore the reduction output for other consumers.
         return is_reduce(x_ins) and x_ins->outputs().size() == 1;
     }
 
@@ -527,7 +525,9 @@ struct find_op_shape_transform_op
                        x_inputs.begin(),
                        reshape_input(x_ins, desc.to_common_from_src()));
         auto new_input_ins = insert(m, x_ins, x_inputs, desc.common_axes_map_from_src());
-        auto new_x_ins     = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
+        auto new_x_ins     = x_ins;
+        if(x_ins->outputs().size() > 1)
+            new_x_ins = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
         if(new_input_ins->get_shape().elements() != input_ins->get_shape().elements())
         {
             auto cdims    = desc.common_dims();
@@ -540,9 +540,12 @@ struct find_op_shape_transform_op
                 return new_input_ins;
             return reshape_input(ins, desc.to_common_from_dst(), true)(input);
         });
-        // Replace old x_ins just in case it is used more than once
-        assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
-        m.replace_instruction(x_ins, new_x_ins);
+        // Restore x_ins for its other consumers.
+        if(new_x_ins != x_ins)
+        {
+            assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
+            m.replace_instruction(x_ins, new_x_ins);
+        }
         // Replace final instruction
         auto pw   = insert(m, ins, inputs, desc.common_axes_map_from_dst());
         auto rins = reshape_input(ins, desc.to_dst_from_common())(pw);
diff --git a/test/simplify_reshapes_test.cpp b/test/simplify_reshapes_test.cpp
index 6d9f59d2081..f274cf37435 100644
--- a/test/simplify_reshapes_test.cpp
+++ b/test/simplify_reshapes_test.cpp
@@ -6494,6 +6494,67 @@ TEST_CASE(op_shape_transform_shadowed_broadcast)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(op_shape_transform_reduce_expanding_chain)
+{
+    migraphx::module m1;
+    {
+        auto image = m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale = m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
+        auto mean = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
+        auto unsqueeze = m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean);
+        auto expanded  = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}), unsqueeze);
+        auto mul = m1.add_instruction(migraphx::make_op("mul"), expanded, scale);
+        m1.add_return({mul});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto image     = m2.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale     = m2.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
+        auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), image);
+        auto expanded  = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 2, 3}}}), unsqueeze);
+        auto mean =
+            m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), expanded);
+        auto mul = m2.add_instruction(migraphx::make_op("mul"), mean, scale);
+        m2.add_return({mul});
+    }
+    EXPECT(m1.sort() == m2.sort());
+}
+
+TEST_CASE(op_shape_transform_shared_reduce_element_preserving_chain)
+{
+    migraphx::module m1;
+    {
+        auto image = m1.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale = m1.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 1, 3}});
+        auto mean = m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
+        auto unsqueeze = m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), mean);
+        auto mul       = m1.add_instruction(migraphx::make_op("mul"), unsqueeze, scale);
+        auto shared    = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), mean);
+        m1.add_return({mul, shared});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto image     = m2.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale     = m2.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 1, 3}});
+        auto unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {3}}}), image);
+        auto mean =
+            m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), unsqueeze);
+        auto squeeze = m2.add_instruction(migraphx::make_op("squeeze", {{"axes", {3}}}), mean);
+        auto mul     = m2.add_instruction(migraphx::make_op("mul"), mean, scale);
+        auto shared  = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), squeeze);
+        m2.add_return({mul, shared});
+    }
+    EXPECT(m1.sort() == m2.sort());
+}
+
 TEST_CASE(op_shape_transform_shared_reduce_expanding_chain)
 {
     migraphx::module m1;

From e1c323c5d3242819f9fe58238d9fc4ab7f15cbb6 Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Wed, 7 Oct 2026 10:48:59 -0500
Subject: [PATCH 05/10] Revert the changes

---
 src/simplify_reshapes.cpp | 25 ++++++++++---------------
 1 file changed, 10 insertions(+), 15 deletions(-)

diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index c1e712783ce..d2eeecb6e0a 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -282,10 +282,9 @@ struct find_op_shape_transform_op
                return not contains({"multibroadcast", "broadcast", "contiguous"}, x->name());
            }))
             return false;
-        if(input_ins->get_shape().elements() == x_ins->get_shape().elements())
-            return true;
-        // Expanding chains cannot restore the reduction output for other consumers.
-        return is_reduce(x_ins) and x_ins->outputs().size() == 1;
+        // A chain that changes the element count is only valid for a reduction
+        return is_reduce(x_ins) or
+               input_ins->get_shape().elements() == x_ins->get_shape().elements();
     }
 
     static bool matches_op(instruction_ref ins)
@@ -525,9 +524,7 @@ struct find_op_shape_transform_op
                        x_inputs.begin(),
                        reshape_input(x_ins, desc.to_common_from_src()));
         auto new_input_ins = insert(m, x_ins, x_inputs, desc.common_axes_map_from_src());
-        auto new_x_ins     = x_ins;
-        if(x_ins->outputs().size() > 1)
-            new_x_ins = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
+        auto new_x_ins     = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
         if(new_input_ins->get_shape().elements() != input_ins->get_shape().elements())
         {
             auto cdims    = desc.common_dims();
@@ -540,12 +537,9 @@ struct find_op_shape_transform_op
                 return new_input_ins;
             return reshape_input(ins, desc.to_common_from_dst(), true)(input);
         });
-        // Restore x_ins for its other consumers.
-        if(new_x_ins != x_ins)
-        {
-            assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
-            m.replace_instruction(x_ins, new_x_ins);
-        }
+        // Replace old x_ins just in case it is used more than once
+        assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
+        m.replace_instruction(x_ins, new_x_ins);
         // Replace final instruction
         auto pw   = insert(m, ins, inputs, desc.common_axes_map_from_dst());
         auto rins = reshape_input(ins, desc.to_dst_from_common())(pw);
@@ -1948,9 +1942,10 @@ struct find_unary_shape_transforms
             else
                 move_down = false;
         }
-        else if(not move_up and not move_down and not yops.empty())
+        else if(not move_up and not move_down)
         {
-            move_up = true;
+            if(not yops.empty())
+                move_up = true;
         }
 
         if(move_up)

From 71c3059ca40ecd5e0f4ab924e79a373eeb47dc1d Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Wed, 7 Oct 2026 12:41:09 -0500
Subject: [PATCH 06/10] Fix src_from_common

---
 .../migraphx/shape_transform_descriptor.hpp   |  2 ++
 src/shape_transform_descriptor.cpp            | 12 +++++++
 src/simplify_reshapes.cpp                     | 32 ++++++++++++++++---
 test/shape_transform_descriptor.cpp           | 28 ++++++++++++++++
 4 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/src/include/migraphx/shape_transform_descriptor.hpp b/src/include/migraphx/shape_transform_descriptor.hpp
index bdd6ab6a5df..d9ff92f5611 100644
--- a/src/include/migraphx/shape_transform_descriptor.hpp
+++ b/src/include/migraphx/shape_transform_descriptor.hpp
@@ -117,6 +117,8 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     shape_transform_descriptor to_common_from_src() const;
     shape_transform_descriptor to_common_from_dst() const;
     shape_transform_descriptor to_dst_from_common() const;
+    // Transform the common dimensions back into the source dimensions. A dimension broadcasted
+    // from no source axis is squeezed away, so it is hidden and must be 1 to rebase.
     shape_transform_descriptor to_src_from_common() const;
 
     // Transform the destination dimensions back into the original dimensions. This is empty when
diff --git a/src/shape_transform_descriptor.cpp b/src/shape_transform_descriptor.cpp
index 6b4e9bf4876..c43a9330c3c 100644
--- a/src/shape_transform_descriptor.cpp
+++ b/src/shape_transform_descriptor.cpp
@@ -769,6 +769,9 @@ shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<
         }
         else if(subs.size() == 1)
         {
+            // A hidden axis of 1 has no broadcast to absorb a different dim
+            if(not broadcast and subs.front()->has_hidden_axis() and subs.front()->len == 1)
+                return {};
             subs.front()->len = dim;
             if(broadcast)
                 subs.front()->hide();
@@ -1990,6 +1993,15 @@ shape_transform_descriptor shape_transform_descriptor::to_src_from_common() cons
         return {subdimensions};
     });
     result.simplify();
+    // A broadcasted dimension with no source axis is squeezed away, so hide
+    // its axis to require it to stay 1 when rebased
+    for_each_subdimension(result.dimensions, [&](dimension::sub& s) {
+        if(s.axis.size() != 1)
+            return;
+        const auto& common_sub = subs[s.axis.front()];
+        if(common_sub.origin_axis().empty() and common_sub.len != 1)
+            s.hide();
+    });
     return result;
 }
 
diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index d2eeecb6e0a..29f199e808c 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -292,6 +292,20 @@ struct find_op_shape_transform_op
         return is_reduce(ins) or ins->get_operator().attributes().contains("pointwise");
     }
 
+    // True when x_ins is only consumed through the chain ending at ins, so
+    // it is dead once ins is rewritten
+    static bool
+    is_private_chain(instruction_ref x_ins, instruction_ref input_ins, instruction_ref ins)
+    {
+        if(not all_of(input_ins->outputs(), [&](instruction_ref out) { return out == ins; }))
+            return false;
+        if(x_ins->outputs().size() != 1)
+            return false;
+        return not any_input_of(input_ins->inputs().front(), x_ins, [](instruction_ref i) {
+            return i->outputs().size() != 1;
+        });
+    }
+
     static bool is_reduce(instruction_ref ins)
     {
         return starts_with(ins->name(), "reduce_") or ins->name() == "argmin" or
@@ -497,6 +511,13 @@ struct find_op_shape_transform_op
         if(not is_valid(x_ins, desc))
             return;
 
+        // x_ins is restored by rebasing the common dims onto its transform, which
+        // fails when the chain broadcasts a new axis, so then it can only be removed
+        auto src_desc        = desc.to_src_from_common();
+        const bool restore_x = not src_desc.rebase(desc.common_dims()).empty();
+        if(not restore_x and not is_private_chain(x_ins, input_ins, ins))
+            return;
+
         // ins is remapped via the dst map; bail if its argmin/argmax axis splits
         if(is_reduce(ins) and ins->get_operator().to_value().contains("axis") and
            not argmax_axis_unsplit(ins, desc.common_axes_map_from_dst()))
@@ -524,7 +545,13 @@ struct find_op_shape_transform_op
                        x_inputs.begin(),
                        reshape_input(x_ins, desc.to_common_from_src()));
         auto new_input_ins = insert(m, x_ins, x_inputs, desc.common_axes_map_from_src());
-        auto new_x_ins     = reshape_input(x_ins, desc.to_src_from_common())(new_input_ins);
+        // Replace old x_ins just in case it is used more than once
+        if(restore_x)
+        {
+            auto new_x_ins = reshape_input(x_ins, src_desc)(new_input_ins);
+            assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
+            m.replace_instruction(x_ins, new_x_ins);
+        }
         if(new_input_ins->get_shape().elements() != input_ins->get_shape().elements())
         {
             auto cdims    = desc.common_dims();
@@ -537,9 +564,6 @@ struct find_op_shape_transform_op
                 return new_input_ins;
             return reshape_input(ins, desc.to_common_from_dst(), true)(input);
         });
-        // Replace old x_ins just in case it is used more than once
-        assert(x_ins->get_shape().lens() == new_x_ins->get_shape().lens());
-        m.replace_instruction(x_ins, new_x_ins);
         // Replace final instruction
         auto pw   = insert(m, ins, inputs, desc.common_axes_map_from_dst());
         auto rins = reshape_input(ins, desc.to_dst_from_common())(pw);
diff --git a/test/shape_transform_descriptor.cpp b/test/shape_transform_descriptor.cpp
index 42af9705b8c..fac00946db8 100644
--- a/test/shape_transform_descriptor.cpp
+++ b/test/shape_transform_descriptor.cpp
@@ -999,6 +999,34 @@ TEST_CASE(common_dims_resize)
            ops{make_op("squeeze", {{"axes", {3, 5}}})});
     EXPECT(desc.to_src_from_common().generate({4, 16, 32, 1, 32, 1}) ==
            ops{make_op("squeeze", {{"axes", {3, 5}}})});
+    // The broadcasted common dims have no source axis, so they cant be rebased back
+    EXPECT(desc.to_src_from_common().rebase({4, 16, 32, 2, 32, 2}).empty());
+}
+
+TEST_CASE(common_dims_reduce_unsqueeze_broadcast)
+{
+    auto desc = make_simple_descriptor({1, 1, 1, 3},
+                                       make_op("unsqueeze", {{"axes", {3}}}),
+                                       make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}))
+                    .rebase({1, 384, 480, 3}, true);
+
+    EXPECT(desc.common_dims() == final_lens{1, 384, 480, 2, 3});
+    EXPECT(desc.find_broadcasted_axes() == std::set<std::size_t>{1, 2});
+
+    EXPECT(desc.common_axes_map_from_src() == axes_map{{0}, {1}, {2}, {4}});
+    EXPECT(desc.common_axes_map_from_dst() == axes_map{{0}, {1}, {2}, {3}, {4}});
+
+    EXPECT(desc.to_common_from_src().generate({1, 384, 480, 3}) ==
+           ops{make_op("unsqueeze", {{"axes", {3}}}),
+               make_op("multibroadcast", {{"out_lens", {1, 384, 480, 2, 3}}})});
+    EXPECT(desc.to_common_from_dst().generate({1, 1, 1, 2, 3}, true) == ops{});
+    EXPECT(desc.to_dst_from_common().generate({1, 1, 1, 2, 3}) == ops{});
+
+    EXPECT(desc.to_src_from_common().generate({1, 1, 1, 1, 3}) ==
+           ops{make_op("squeeze", {{"axes", {3}}})});
+    // The broadcasted common dim has no source axis, so it cant be rebased back
+    EXPECT(desc.to_src_from_common().rebase({1, 1, 1, 2, 3}).empty());
+    EXPECT(desc.to_src_from_common().rebase(desc.common_dims()).empty());
 }
 
 TEST_CASE(common_dims_squeeze_1x1)

From 9db7e41a70ce37a8f47e2c16a240cc395bba4470 Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Wed, 7 Oct 2026 15:22:30 -0500
Subject: [PATCH 07/10] Fix the
 (op_shape_transform_shared_reduce_expanding_chain test case

---
 .../migraphx/shape_transform_descriptor.hpp   |  2 +
 src/shape_transform_descriptor.cpp            | 79 +++++++++++++++++--
 src/simplify_reshapes.cpp                     | 36 ++++++---
 test/fuse_reduce.cpp                          |  2 +-
 test/rewrite_reduce.cpp                       | 10 +--
 test/shape_transform_descriptor.cpp           | 60 +++++++++++++-
 test/simplify_reshapes_test.cpp               | 26 ++++--
 7 files changed, 178 insertions(+), 37 deletions(-)

diff --git a/src/include/migraphx/shape_transform_descriptor.hpp b/src/include/migraphx/shape_transform_descriptor.hpp
index d9ff92f5611..ad48bb7a743 100644
--- a/src/include/migraphx/shape_transform_descriptor.hpp
+++ b/src/include/migraphx/shape_transform_descriptor.hpp
@@ -104,6 +104,8 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
                                     bool no_broadcast                          = false) const;
 
     std::set<std::size_t> find_broadcasted_axes() const;
+    // Returns true if a dimension is broadcasted without a source axis, ie a new dimension
+    bool has_axisless_broadcast() const;
     bool has_broadcast() const;
     // Returns true if any non-trivial axes are reordered, requiring a transpose
     bool is_transposed() const;
diff --git a/src/shape_transform_descriptor.cpp b/src/shape_transform_descriptor.cpp
index c43a9330c3c..aa7267fb75e 100644
--- a/src/shape_transform_descriptor.cpp
+++ b/src/shape_transform_descriptor.cpp
@@ -215,6 +215,23 @@ static std::size_t visible_len(const std::vector<dimension::sub*>& subs)
         });
 }
 
+static bool has_no_axes(const dimension& d)
+{
+    return std::all_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
+        return s.axis.empty() and s.hidden_axis.empty();
+    });
+}
+
+// Length of the dimension before any broadcast is applied
+static std::size_t input_len(const dimension& d)
+{
+    return transform_accumulate(d.subdimensions.begin(),
+                                d.subdimensions.end(),
+                                std::size_t{1},
+                                std::multiplies<>{},
+                                [](const dimension::sub& s) { return s.axis.empty() ? 1 : s.len; });
+}
+
 static std::vector<std::size_t> compute_dims(const operation& op,
                                              const std::vector<std::size_t>& idims)
 {
@@ -742,10 +759,49 @@ struct rebase_ambiguity_resolver
     static const std::size_t last_axis_split = std::numeric_limits<std::size_t>::max();
 };
 
+// A hidden broadcast of an axis that is rebased to another size cant stay
+// attached to that axis. When a dimension of 1 without an axis precedes it,
+// the axes shift down into that dimension and the broadcast becomes its own
+// axis-less dimension, as if the axis had been unsqueezed and broadcasted.
+static void detach_unabsorbed_broadcasts(std::vector<dimension>& dimensions,
+                                         const std::vector<std::size_t>& dims)
+{
+    auto axes_map = group_axes(dimensions);
+    for(std::size_t h : range(dimensions.size()))
+    {
+        auto& subs = dimensions[h].subdimensions;
+        if(subs.size() != 1)
+            continue;
+        auto& s = subs.front();
+        if(not s.has_hidden_axis() or s.hidden_axis.size() != 1)
+            continue;
+        auto axis = s.hidden_axis.front();
+        assert(axis < dims.size());
+        if(axes_map.at(axis).size() != 1)
+            continue;
+        auto dim = dims[axis];
+        if(dim == 0 or dim == 1 or dim == s.len)
+            continue;
+        // Search backwards through dimensions of 1 for one without an axis
+        auto rit = std::find_if(std::make_reverse_iterator(dimensions.begin() + h),
+                                dimensions.rend(),
+                                [](const dimension& d) { return d.len() != 1 or has_no_axes(d); });
+        if(rit == dimensions.rend() or rit->len() != 1)
+            continue;
+        dimension bdim{{dimension::sub{s.len, {}}}};
+        s.expose();
+        s.len = 1;
+        dimensions.erase(std::prev(rit.base()));
+        dimensions.insert(dimensions.begin() + h, bdim);
+    }
+}
+
 shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<std::size_t>& dims,
                                                               bool broadcast) const
 {
-    auto result   = *this;
+    auto result = *this;
+    if(broadcast)
+        detach_unabsorbed_broadcasts(result.dimensions, dims);
     auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
     for(auto& [axis, subs] : axes_map)
     {
@@ -1707,12 +1763,6 @@ struct operation_list
 
 } // namespace
 
-static bool has_no_axes(const dimension& d)
-{
-    return std::all_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
-        return s.axis.empty() and s.hidden_axis.empty();
-    });
-}
 static bool has_axes(const dimension& d)
 {
     return std::any_of(d.subdimensions.begin(), d.subdimensions.end(), [](const dimension::sub& s) {
@@ -1796,6 +1846,14 @@ shape_transform_descriptor::generate(const std::vector<std::size_t>& input_dims,
             else
             {
                 result.push_back(make_op("multibroadcast", {{"out_lens", out_lens}}));
+                // Broadcasting aligns from the right, so a dimension without an
+                // axis that only follows input dimensions of 1 needs no unsqueeze
+                auto first_non1 =
+                    std::find_if(new_dims.begin(), new_dims.end(), [](const dimension& d) {
+                        return input_len(d) != 1;
+                    });
+                new_dims.erase(std::remove_if(new_dims.begin(), first_non1, &has_no_axes),
+                               first_non1);
             }
         }
         // If all the dimensions have no axes then there isnt anthing else to do
@@ -1863,6 +1921,13 @@ std::set<std::size_t> shape_transform_descriptor::find_broadcasted_axes() const
     return result;
 }
 
+bool shape_transform_descriptor::has_axisless_broadcast() const
+{
+    auto subs = get_all_subdimensions(dimensions);
+    return std::any_of(subs.begin(), subs.end(), [](const dimension::sub& s) {
+        return s.origin_axis().empty() and s.len != 1;
+    });
+}
 bool shape_transform_descriptor::has_broadcast() const
 {
     return std::any_of(dimensions.begin(), dimensions.end(), [&](const dimension& d) {
diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index 29f199e808c..8a1b718ed83 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -277,11 +277,17 @@ struct find_op_shape_transform_op
         // shape_transform_descriptor doesnt handle scalars for now
         if(input_ins->get_shape().scalar() or x_ins->get_shape().scalar())
             return false;
-        // If its just a broadcast then skip
+        // A chain of only broadcasts is skipped, unless it broadcasts a new
+        // axis of a reduction, which is the same as an unsqueeze and broadcast
         if(not any_input_of(input_ins, x_ins, [](instruction_ref x) {
                return not contains({"multibroadcast", "broadcast", "contiguous"}, x->name());
            }))
-            return false;
+        {
+            if(not is_reduce(x_ins))
+                return false;
+            auto desc = make_descriptor(x_ins, chain_ops(x_ins, input_ins), input_ins);
+            return not desc.empty() and desc.has_axisless_broadcast();
+        }
         // A chain that changes the element count is only valid for a reduction
         return is_reduce(x_ins) or
                input_ins->get_shape().elements() == x_ins->get_shape().elements();
@@ -312,6 +318,20 @@ struct find_op_shape_transform_op
                ins->name() == "argmax";
     }
 
+    // The shape transform operators from x_ins to input_ins in order
+    static std::vector<operation> chain_ops(instruction_ref x_ins, instruction_ref input_ins)
+    {
+        std::vector<operation> ops;
+        auto next_ins = input_ins;
+        while(next_ins != x_ins)
+        {
+            ops.push_back(next_ins->get_operator());
+            next_ins = next_ins->inputs().front();
+        }
+        std::reverse(ops.begin(), ops.end());
+        return ops;
+    }
+
     template <class F>
     static instruction_ref find_input_if(instruction_ref start, instruction_ref last, F f)
     {
@@ -494,17 +514,7 @@ struct find_op_shape_transform_op
         auto x_ins     = r.instructions["x"];
         auto input_ins = r.instructions["input"];
 
-        std::vector<operation> ops;
-        auto next_ins = input_ins;
-        while(next_ins != x_ins)
-        {
-            ops.push_back(next_ins->get_operator());
-            next_ins = next_ins->inputs().front();
-        }
-        assert(next_ins == x_ins);
-        std::reverse(ops.begin(), ops.end());
-
-        auto desc = make_descriptor(x_ins, ops, input_ins);
+        auto desc = make_descriptor(x_ins, chain_ops(x_ins, input_ins), input_ins);
         if(desc.empty())
             return;
 
diff --git a/test/fuse_reduce.cpp b/test/fuse_reduce.cpp
index f305de52bfc..2d43fca990d 100644
--- a/test/fuse_reduce.cpp
+++ b/test/fuse_reduce.cpp
@@ -1978,7 +1978,7 @@ TEST_CASE(pointwise_reshapes_reduce_shadowed_broadcast)
         auto w   = mm->add_parameter("w", ws);
         auto pw0 = add_pointwise(p2, "main:pointwise0", {x}, single_pointwise("sqrt"));
         auto wu  = mm->add_instruction(migraphx::make_op("unsqueeze", {{"axes", {0}}}), w);
-        auto xr  = mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 1, 2, 4}}}), pw0);
+        auto xr  = mm->add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 4}}}), pw0);
         auto xb  = mm->add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 2, 4}}}), xr);
         auto rsum = add_reduce(
diff --git a/test/rewrite_reduce.cpp b/test/rewrite_reduce.cpp
index 2923fb439a5..76190bcb9bc 100644
--- a/test/rewrite_reduce.cpp
+++ b/test/rewrite_reduce.cpp
@@ -114,12 +114,11 @@ TEST_CASE(dot_skinny_rewrite)
     {
         auto a       = m2.add_parameter("a", a_shape);
         auto b       = m2.add_parameter("b", b_shape);
-        auto a_unsq  = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), a);
         auto b_unsq  = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), b);
         auto b_trans = m2.add_instruction(
             migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), b_unsq);
-        auto a_bc = m2.add_instruction(
-            migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 128}}}), a_unsq);
+        auto a_bc =
+            m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {1, 4, 128}}}), a);
         auto mul = m2.add_instruction(migraphx::make_op("mul"), a_bc, b_trans);
         auto red = m2.add_instruction(migraphx::make_op("reduce_sum", {{"axes", {2}}}), mul);
         auto sq  = m2.add_instruction(migraphx::make_op("squeeze", {{"axes", {2}}}), red);
@@ -396,9 +395,8 @@ TEST_CASE(dot_softmax_dot_projection_rewrite)
             m2.add_instruction(migraphx::make_op("reshape", {{"dims", {1, 2, 1, 16}}}), trans);
         auto a_bc = m2.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 4, 16}}}), rsp);
-        auto unsq_w  = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), w);
-        auto trans_w = m2.add_instruction(
-            migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), unsq_w);
+        auto trans_w =
+            m2.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 1}}}), w);
         auto w_bc = m2.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 2, 4, 16}}}), trans_w);
         auto mul = m2.add_instruction(migraphx::make_op("mul"), a_bc, w_bc);
diff --git a/test/shape_transform_descriptor.cpp b/test/shape_transform_descriptor.cpp
index fac00946db8..ea119c9a288 100644
--- a/test/shape_transform_descriptor.cpp
+++ b/test/shape_transform_descriptor.cpp
@@ -654,12 +654,33 @@ TEST_CASE(optimize_multibroadcast_transpose)
                    make_op("multibroadcast", {{"out_lens", {2, 320, 64, 64}}}),
                    make_op("transpose", {{"permutation", {0, 2, 3, 1}}}),
                }) == ops{
-                         make_op("unsqueeze", {{"axes", {0}}}),
-                         make_op("transpose", {{"permutation", {0, 2, 3, 1}}}),
+                         make_op("transpose", {{"permutation", {1, 2, 0}}}),
                          make_op("multibroadcast", {{"out_lens", {2, 64, 64, 320}}}),
                      });
 }
 
+TEST_CASE(optimize_unsqueeze_multibroadcast_after_1s)
+{
+    // Broadcasting aligns from the right, so an unsqueezed axis that only
+    // follows 1s is added by the multibroadcast
+    EXPECT(check_optimize_shape_transforms(
+               {1, 1, 1, 3},
+               {
+                   make_op("unsqueeze", {{"axes", {3}}}),
+                   make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}),
+               }) == ops{make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}})});
+    // A non-1 axis before the unsqueezed axis still needs the unsqueeze
+    EXPECT(check_optimize_shape_transforms(
+               {2, 1, 3},
+               {
+                   make_op("unsqueeze", {{"axes", {2}}}),
+                   make_op("multibroadcast", {{"out_lens", {2, 1, 4, 3}}}),
+               }) == ops{
+                         make_op("unsqueeze", {{"axes", {2}}}),
+                         make_op("multibroadcast", {{"out_lens", {2, 1, 4, 3}}}),
+                     });
+}
+
 TEST_CASE(optimize_unsqueeze_transpose_squeeze_multibroadcast)
 {
     EXPECT(check_optimize_shape_transforms(
@@ -1025,10 +1046,43 @@ TEST_CASE(common_dims_reduce_unsqueeze_broadcast)
     EXPECT(desc.to_src_from_common().generate({1, 1, 1, 1, 3}) ==
            ops{make_op("squeeze", {{"axes", {3}}})});
     // The broadcasted common dim has no source axis, so it cant be rebased back
+    EXPECT(desc.has_axisless_broadcast());
     EXPECT(desc.to_src_from_common().rebase({1, 1, 1, 2, 3}).empty());
     EXPECT(desc.to_src_from_common().rebase(desc.common_dims()).empty());
 }
 
+TEST_CASE(rebase_broadcast_detach_new_axis)
+{
+    // The broadcast is attributed to source axis 2 while it is 1, but once
+    // that axis is rebased to 480 the broadcast becomes a new axis
+    auto base_desc = make_simple_descriptor(
+        {1, 1, 1, 3}, make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}));
+    EXPECT(not base_desc.has_axisless_broadcast());
+    EXPECT(base_desc.find_broadcasted_axes() == std::set<std::size_t>{2});
+
+    auto desc = base_desc.rebase({1, 384, 480, 3}, true);
+    EXPECT(desc.common_dims() == final_lens{1, 384, 480, 2, 3});
+    EXPECT(desc.find_broadcasted_axes() == std::set<std::size_t>{1, 2});
+    EXPECT(desc.has_axisless_broadcast());
+    EXPECT(desc.common_axes_map_from_src() == axes_map{{0}, {1}, {2}, {4}});
+    EXPECT(desc.to_common_from_src().generate({1, 384, 480, 3}) ==
+           ops{make_op("unsqueeze", {{"axes", {3}}}),
+               make_op("multibroadcast", {{"out_lens", {1, 384, 480, 2, 3}}})});
+    EXPECT(desc.to_src_from_common().rebase(desc.common_dims()).empty());
+}
+
+TEST_CASE(rebase_broadcast_reduced_axis)
+{
+    // Broadcasting a reduced axis back to its size stays attached to that axis
+    auto desc = make_simple_descriptor({2, 32, 1},
+                                       make_op("multibroadcast", {{"out_lens", {2, 32, 40960}}}))
+                    .rebase({2, 32, 40960}, true);
+    EXPECT(desc.common_dims() == final_lens{2, 32, 40960});
+    EXPECT(desc.find_broadcasted_axes() == std::set<std::size_t>{2});
+    EXPECT(not desc.has_axisless_broadcast());
+    EXPECT(desc.to_src_from_common().generate({2, 32, 1}) == ops{});
+}
+
 TEST_CASE(common_dims_squeeze_1x1)
 {
     auto desc = make_simple_descriptor({1, 1}, make_op("squeeze", {{"axes", {0}}}));
@@ -1494,7 +1548,7 @@ TEST_CASE(generate_shape_transforms_for)
 
     EXPECT(generate_for({2, 2, 2, 2, 3}, {0, 2, 0, 1, 0}, {4}) ==
            ops{
-               make_op("reshape", {{"dims", {1, 2, 1, 2, 1}}}),
+               make_op("reshape", {{"dims", {2, 1, 2, 1}}}),
                make_op("multibroadcast", {{"out_lens", {2, 2, 2, 2, 3}}}),
            });
 
diff --git a/test/simplify_reshapes_test.cpp b/test/simplify_reshapes_test.cpp
index f274cf37435..bc28b24cfd7 100644
--- a/test/simplify_reshapes_test.cpp
+++ b/test/simplify_reshapes_test.cpp
@@ -5006,7 +5006,7 @@ TEST_CASE(conv_add_layernorm_conv)
             reshape2);
         auto add1           = m2.add_instruction(migraphx::make_op("add"), reshape1, reshape2b);
         auto unsqueeze_p_y2 = m2.add_instruction(
-            migraphx::make_op("unsqueeze", {{"axes", {1, 2, 3, 4}}, {"steps", {}}}), p_y2);
+            migraphx::make_op("unsqueeze", {{"axes", {1, 2}}, {"steps", {}}}), p_y2);
         auto unsqueeze_p_y2b = m2.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {2, 32, 10, 64, 64}}}),
             unsqueeze_p_y2);
@@ -5018,7 +5018,7 @@ TEST_CASE(conv_add_layernorm_conv)
         auto sub1           = m2.add_instruction(migraphx::make_op("sub"), add1, reduce_sum1b);
         auto mul1           = m2.add_instruction(migraphx::make_op("mul"), add1, add1);
         auto unsqueeze_p_y3 = m2.add_instruction(
-            migraphx::make_op("unsqueeze", {{"axes", {1, 2, 3, 4}}, {"steps", {}}}), p_y3);
+            migraphx::make_op("unsqueeze", {{"axes", {1, 2}}, {"steps", {}}}), p_y3);
         auto p_y3b = m2.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {2, 32, 10, 64, 64}}}),
             unsqueeze_p_y3);
@@ -5546,7 +5546,7 @@ TEST_CASE(slice_reshape_multibroadcast_rebase_axis)
     }
     auto m2 = m1;
     run_pass(m1);
-    EXPECT(m1.get_output_shapes() == m2.get_output_shapes());
+    EXPECT(to_lens(m1.get_output_shapes()) == to_lens(m2.get_output_shapes()));
 }
 
 // The gather_slice_concat tests below use parameter indices (rather than
@@ -6431,8 +6431,7 @@ TEST_CASE(dequantizelinear_entry_shape_transform_zero_point)
         auto scale_bcast = m2.add_instruction(
             migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", {4, 3, 2}}}), scale_reshape);
         auto x_reshape = m2.add_instruction(migraphx::make_op("reshape", {{"dims", {4, 3, 2}}}), x);
-        auto zp_unsqueeze =
-            m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1, 2}}}), zp);
+        auto zp_unsqueeze = m2.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {1}}}), zp);
         auto zp_bcast = m2.add_instruction(
             migraphx::make_op("multibroadcast", {{"out_lens", {4, 3, 2}}}), zp_unsqueeze);
         auto dq = m2.add_instruction(
@@ -6570,9 +6569,22 @@ TEST_CASE(op_shape_transform_shared_reduce_expanding_chain)
             migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), mean);
         m1.add_return({mul, shared});
     }
-    auto m2 = m1;
     run_pass(m1);
-    EXPECT(m1 == m2);
+
+    // The shared reduce cant be moved, so only the unsqueeze folds into the broadcast
+    migraphx::module m2;
+    {
+        auto image = m2.add_parameter("image", {migraphx::shape::float_type, {1, 384, 480, 3}});
+        auto scale = m2.add_parameter("scale", {migraphx::shape::float_type, {1, 1, 1, 2, 3}});
+        auto mean = m2.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {1, 2}}}), image);
+        auto expanded = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 1, 1, 2, 3}}}), mean);
+        auto mul    = m2.add_instruction(migraphx::make_op("mul"), expanded, scale);
+        auto shared = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {1, 384, 480, 3}}}), mean);
+        m2.add_return({mul, shared});
+    }
+    EXPECT(m1.sort() == m2.sort());
 }
 
 int main(int argc, const char* argv[]) { test::run(argc, argv); }

From c31f2ead507d81509fb9c2afda426a2eaf781060 Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Fri, 9 Oct 2026 13:20:17 -0500
Subject: [PATCH 08/10] Fix tidy warnings

---
 src/shape_transform_descriptor.cpp | 88 ++++++++++++++++--------------
 1 file changed, 48 insertions(+), 40 deletions(-)

diff --git a/src/shape_transform_descriptor.cpp b/src/shape_transform_descriptor.cpp
index aa7267fb75e..a7d580795f9 100644
--- a/src/shape_transform_descriptor.cpp
+++ b/src/shape_transform_descriptor.cpp
@@ -796,58 +796,66 @@ static void detach_unabsorbed_broadcasts(std::vector<dimension>& dimensions,
     }
 }
 
-shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<std::size_t>& dims,
-                                                              bool broadcast) const
+// Adjust the subdimensions of an axis to its new dim. Returns false when the
+// axis cant be rebased.
+static bool rebase_axis(const std::vector<dimension::sub*>& subs, std::size_t dim, bool broadcast)
 {
-    auto result = *this;
-    if(broadcast)
-        detach_unabsorbed_broadcasts(result.dimensions, dims);
-    auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
-    for(auto& [axis, subs] : axes_map)
+    if(dim == len(subs))
     {
-        assert(axis < dims.size());
-        auto dim       = dims[axis];
-        if(dim == len(subs))
-        {
-            if(not broadcast)
-            {
-                for(auto* sub : subs)
-                    sub->expose();
-            }
-        }
-        else if(dim == 1)
+        if(not broadcast)
         {
             for(auto* sub : subs)
-            {
-                if(not sub->has_hidden_axis())
-                    sub->len = 1;
-            }
+                sub->expose();
         }
-        else if(subs.size() == 1)
+    }
+    else if(dim == 1)
+    {
+        for(auto* sub : subs)
         {
-            // A hidden axis of 1 has no broadcast to absorb a different dim
-            if(not broadcast and subs.front()->has_hidden_axis() and subs.front()->len == 1)
-                return {};
-            subs.front()->len = dim;
-            if(broadcast)
-                subs.front()->hide();
-            else
-                subs.front()->expose();
+            if(not sub->has_hidden_axis())
+                sub->len = 1;
         }
-        else if(dim == visible_len(subs))
+    }
+    else if(subs.size() == 1)
+    {
+        // A hidden axis of 1 has no broadcast to absorb a different dim
+        if(not broadcast and subs.front()->has_hidden_axis() and subs.front()->len == 1)
+            return false;
+        subs.front()->len = dim;
+        if(broadcast)
+            subs.front()->hide();
+        else
+            subs.front()->expose();
+    }
+    else if(dim == visible_len(subs))
+    {
+        for(auto* sub : subs)
         {
-            for(auto* sub : subs)
+            if(sub->has_hidden_axis())
             {
-                if(sub->has_hidden_axis())
-                {
-                    sub->expose();
-                    sub->len = 1;
-                }
+                sub->expose();
+                sub->len = 1;
             }
         }
-        else
-            return {};
     }
+    else
+        return false;
+    return true;
+}
+
+shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<std::size_t>& dims,
+                                                              bool broadcast) const
+{
+    auto result = *this;
+    if(broadcast)
+        detach_unabsorbed_broadcasts(result.dimensions, dims);
+    auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
+    if(not std::all_of(axes_map.begin(), axes_map.end(), [&](auto& p) {
+           const auto& [axis, subs] = p;
+           assert(axis < dims.size());
+           return rebase_axis(subs, dims[axis], broadcast);
+       }))
+        return {};
     for(auto& dim : result.dimensions)
         remove_empty_sub_dims(dim.subdimensions);
     if(broadcast and not is_broadcast_only(dimensions, result.dimensions))

From b187b53959bd63d909af3205a9befe25ab692333 Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Fri, 9 Oct 2026 18:57:24 -0500
Subject: [PATCH 09/10] Fix tidy and cppcheck issues

---
 src/shape_transform_descriptor.cpp | 2 +-
 src/simplify_reshapes.cpp          | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/shape_transform_descriptor.cpp b/src/shape_transform_descriptor.cpp
index a7d580795f9..0420e548615 100644
--- a/src/shape_transform_descriptor.cpp
+++ b/src/shape_transform_descriptor.cpp
@@ -850,7 +850,7 @@ shape_transform_descriptor shape_transform_descriptor::rebase(const std::vector<
     if(broadcast)
         detach_unabsorbed_broadcasts(result.dimensions, dims);
     auto axes_map = rebase_ambiguity_resolver{result, dims}.resolve();
-    if(not std::all_of(axes_map.begin(), axes_map.end(), [&](auto& p) {
+    if(not std::all_of(axes_map.begin(), axes_map.end(), [&](const auto& p) {
            const auto& [axis, subs] = p;
            assert(axis < dims.size());
            return rebase_axis(subs, dims[axis], broadcast);
diff --git a/src/simplify_reshapes.cpp b/src/simplify_reshapes.cpp
index 8a1b718ed83..d626559e3f8 100644
--- a/src/simplify_reshapes.cpp
+++ b/src/simplify_reshapes.cpp
@@ -1976,10 +1976,9 @@ struct find_unary_shape_transforms
             else
                 move_down = false;
         }
-        else if(not move_up and not move_down)
+        else if(not move_up and not move_down and not yops.empty())
         {
-            if(not yops.empty())
-                move_up = true;
+            move_up = true;
         }
 
         if(move_up)

From 085db1094e6dd53edd53aba4de08fbaa3535f7c8 Mon Sep 17 00:00:00 2001
From: Paul <pfultz2@yahoo.com>
Date: Sat, 10 Oct 2026 17:35:04 -0500
Subject: [PATCH 10/10] Fix ASAN error

---
 src/simplify_algebra.cpp       | 14 +++++++++-----
 test/simplify_algebra_test.cpp | 27 +++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/simplify_algebra.cpp b/src/simplify_algebra.cpp
index 22c7e0e03de..1191d663ca2 100644
--- a/src/simplify_algebra.cpp
+++ b/src/simplify_algebra.cpp
@@ -1019,13 +1019,17 @@ struct find_concat_op
         }
         else if(op.name() == "multibroadcast")
         {
-            shape bshape = (*start)->get_shape();
-            auto input   = (*start)->inputs()[0];
-            if(iaxis >= bshape.strides().size() or bshape.strides()[iaxis] == 0)
+            auto input_rank = (*start)->inputs().front()->get_shape().ndim();
+            // The inputs are concatenated, so the axis cant be broadcasted and
+            // the inputs need the same rank for it to map to the same input axis
+            if(std::any_of(start, last, [&](instruction_ref i) {
+                   const auto& s = i->get_shape();
+                   return iaxis >= s.ndim() or s.strides()[iaxis] == 0 or
+                          i->inputs().front()->get_shape().ndim() != input_rank;
+               }))
                 return {start, last};
             op.from_value({{"out_lens", get_output_lens(start, last, iaxis)}});
-            auto delta = bshape.lens().size() - input->get_shape().lens().size();
-            iaxis -= delta;
+            iaxis -= (*start)->get_shape().ndim() - input_rank;
         }
         else if(op.name() == "unsqueeze")
         {
diff --git a/test/simplify_algebra_test.cpp b/test/simplify_algebra_test.cpp
index 485403613b2..2583ac73497 100644
--- a/test/simplify_algebra_test.cpp
+++ b/test/simplify_algebra_test.cpp
@@ -1315,6 +1315,33 @@ TEST_CASE(simplify_concat_mul_broadcast_diff_size)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(simplify_concat_multibroadcast_diff_rank)
+{
+    auto s = migraphx::shape{migraphx::shape::int32_type, {4, 2, 2, 1}};
+    migraphx::module m1;
+    {
+        auto x    = m1.add_parameter("x", s);
+        auto lit1 = m1.add_literal(
+            migraphx::generate_literal({migraphx::shape::int32_type, {4, 1, 1, 1}}, 1));
+        auto lit2 =
+            m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {2, 1, 1}}, 2));
+        auto lit3 =
+            m1.add_literal(migraphx::generate_literal({migraphx::shape::int32_type, {2, 1}}, 3));
+        auto lit1b =
+            m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s.lens()}}), lit1);
+        auto lit2b =
+            m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s.lens()}}), lit2);
+        auto lit3b =
+            m1.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", s.lens()}}), lit3);
+        auto concat =
+            m1.add_instruction(migraphx::make_op("concat", {{"axis", 3}}), lit1b, lit2b, lit3b, x);
+        m1.add_return({concat});
+    }
+    migraphx::module m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(simplify_concat_clip)
 {
     auto s = migraphx::shape{migraphx::shape::int32_type, {1}};