Merge pull request #29829 from spmallick:codex/einsum-dynamic-shape-reset

DNN: reset Einsum shape state when operand dimensions change 🤖🤖🤖
This commit is contained in:
Alexander Smorkalov
2026-08-31 08:09:47 +03:00
committed by GitHub
2 changed files with 53 additions and 19 deletions

View File

@@ -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();

View File

@@ -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);
}
}
}}