#include <cmath>
+#include <limits>
#include "nnc/core/IR/model/actions/ShapeInference.h"
{
// Reshape should have it's output shape filled by importer/user
fillInputShapes(node, op);
+ auto& inShape = op.getInputShape(0);
+ auto outShape = op.getOutputShape(0);
+
+ //-1 in shape underflows to UINT32_MAX
+ const auto autoDimension = std::numeric_limits<uint32_t>::max();
+
+ assert(inShape.rank() == outShape.rank());
+
+ auto inElementsNum = num_elements(inShape);
+ uint32_t outElementsNum = 1;
+ //can't use num_elements due to -1 in input shape and Shape using unsigned ints for dimensions
+ for( uint32_t d = 0; d < outShape.rank(); ++d ) {
+ auto dim = outShape.dim(d);
+ if( dim != autoDimension ) {
+ outElementsNum *= dim;
+ }
+ }
+
+ for( uint32_t d = 0; d < outShape.rank(); ++d ) {
+ auto& dim = outShape.dim(d);
+ if( dim == autoDimension ) {
+ dim = static_cast<uint32_t>(inElementsNum / outElementsNum);
+ }
+ }
+
+ op.setOutputShape(0, outShape);
}
} // namespace model
--- /dev/null
+#include "nnc/core/IR/model/graph/graph.h"
+#include "nnc/core/IR/model/actions/ShapeInference.h"
+#include "nnc/core/IR/model/operations/reshape_op.h"
+#include "nncc/core/ADT/tensor/Shape.h"
+
+#include "gtest/gtest.h"
+
+using namespace nncc::contrib::core::IR::model;
+using nncc::core::ADT::tensor::Shape;
+
+TEST(ShapeInferenceTest, ReshapeAutoDimension) {
+ Graph g;
+ ShapeInference si;
+
+ Shape resultShape{ 10, 1, 10 };
+
+ auto input = g.create<ops::VariableOp>("input");
+ input->getOperation()->setOutputShape(0, Shape{ 10, 2, 5} );
+
+ auto n = g.create<ops::ReshapeOp>("reshape");
+ n->getOperation()->setInputShape( 0, Shape{10, 2, 5} );
+ n->getOperation()->setOutputShape(0, Shape{10, 1, static_cast<uint32_t >(-1)} );
+ n->connectInputTo(0, input->getOutput(0));
+
+ si.visit(n, *static_cast<ops::ReshapeOp*>(n->getOperation()));
+
+ ASSERT_EQ(resultShape, n->getOperation()->getOutputShape(0));
+}