From e9f1117b2f7e5edb7afc14f219758647b837caa1 Mon Sep 17 00:00:00 2001
From: Satya Mallick <spmallick@learnopencv.com>
Date: Sun, 30 Aug 2026 11:30:50 +0530
Subject: [PATCH] DNN: reset Einsum state after operand shape changes

---
 modules/dnn/src/layers/einsum_layer.cpp | 47 +++++++++++++++----------
 modules/dnn/test/test_layers_1d.cpp     | 25 +++++++++++++
 2 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/modules/dnn/src/layers/einsum_layer.cpp b/modules/dnn/src/layers/einsum_layer.cpp
index 4699fa5f81..b5e973c9f3 100644
--- a/modules/dnn/src/layers/einsum_layer.cpp
+++ b/modules/dnn/src/layers/einsum_layer.cpp
@@ -380,6 +380,27 @@ public:
     mutable bool outputShapeComputed;
     mutable MatShape cachedOutputShape;
 
+    // Equation parsing maps labels to concrete dimensions, so every derived
+    // table must be rebuilt together when any operand shape changes.
+    void resetShapeState()
+    {
+        einsumInpShapes.clear();
+        preProcessedInputs.clear();
+        homogenizedInputDims.clear();
+        einsumOutDims.clear();
+        inputSubscriptIndices.clear();
+        subscriptIndicesToLastInput.clear();
+        subscriptIndicesToDimValue.clear();
+        subscriptIndicesToOutputIndices.clear();
+        letter2count.fill(0);
+        letter2index.fill(-1);
+        numLetterIndices = 0;
+        numOfEllipsisDims = 0;
+        numInputs = 0;
+        cachedOutputShape.clear();
+        outputShapeComputed = false;
+    }
+
     void parseEquation(String equation);
     void processEquation(const std::vector<MatShape>& inputs);
     void processBroadcastedDims();
@@ -404,12 +425,13 @@ public:
     );
 
     void computeOutputShape(const std::vector<MatShape>& inputs) const {
-        if (!outputShapeComputed) {
-            // Copy of the existing computation logic
-            const_cast<LayerEinsumImpl*>(this)->processEquation(inputs);
-            const_cast<LayerEinsumImpl*>(this)->processBroadcastedDims();
-            const_cast<LayerEinsumImpl*>(this)->validateOutputSubscript();
-            const_cast<LayerEinsumImpl*>(this)->calculateOutputShape();
+        if (!outputShapeComputed || inputs != einsumInpShapes) {
+            LayerEinsumImpl* self = const_cast<LayerEinsumImpl*>(this);
+            self->resetShapeState();
+            self->processEquation(inputs);
+            self->processBroadcastedDims();
+            self->validateOutputSubscript();
+            self->calculateOutputShape();
 
             cachedOutputShape = einsumOutDims;
             outputShapeComputed = true;
@@ -472,19 +494,6 @@ public:
         CV_UNUSED(requiredOutputs);
         CV_UNUSED(internals);
 
-        // check if input einsumInputShapes is empty
-        if (einsumInpShapes.empty()) {
-            outputShapeComputed = false;
-        } else {
-            // check weather shapes in inputs are compatible with shapes in einsumInpShapes
-            for (int i = 0; i < inputs.size(); i++) {
-                if (inputs[i] != einsumInpShapes[i]) {
-                    outputShapeComputed = false;
-                    break;
-                }
-            }
-        }
-
         computeOutputShape(inputs);
 
         outputs.clear();
diff --git a/modules/dnn/test/test_layers_1d.cpp b/modules/dnn/test/test_layers_1d.cpp
index c5b8957544..bce09bd282 100644
--- a/modules/dnn/test/test_layers_1d.cpp
+++ b/modules/dnn/test/test_layers_1d.cpp
@@ -1659,5 +1659,30 @@ INSTANTIATE_TEST_CASE_P(/*nothing*/, Layer_Einsum_Test, testing::Values(
     std::make_tuple(std::vector<int>({4, 4}), std::vector<int>({4, 4}), "ij,ij->i")
     ));
 
+TEST(Layer_Einsum, DynamicLeadingDimensionReusesLayer)
+{
+    LayerParams lp;
+    lp.type = "Einsum";
+    lp.name = "dynamic_leading_dimension";
+    lp.set("equation", "mc,mchw->mhw");
+    Ptr<Layer> layer = EinsumLayer::create(lp);
+
+    const int channels = 3, height = 2, width = 2;
+    const int leadingDimensions[] = {2, 5, 2};
+    for (const int leadingDimension : leadingDimensions)
+    {
+        Mat coefficients(MatShape{leadingDimension, channels}, CV_32F, Scalar::all(1));
+        Mat features(MatShape{leadingDimension, channels, height, width}, CV_32F, Scalar::all(2));
+        std::vector<Mat> inputs{coefficients, features}, outputs;
+
+        runLayer(layer, inputs, outputs);
+
+        ASSERT_EQ(outputs.size(), (size_t)1);
+        EXPECT_EQ(shape(outputs[0]), MatShape({leadingDimension, height, width}));
+        Mat expected(MatShape{leadingDimension, height, width}, CV_32F, Scalar::all(6));
+        normAssert(outputs[0], expected);
+    }
+}
+
 
 }}