bool fix_shape(loco::ConstGen *node)
{
- // TODO fill this
- return false;
+ const ShapeInferenceData *shapedata = node->annot<ShapeInferenceData>();
+ if (shapedata != nullptr)
+ {
+ // shape inference is already done for ConstGen
+ return false;
+ }
+
+ // ConstGen itself has shape information, copy them
+ auto shape_data = stdex::make_unique<ShapeInferenceData>();
+ uint32_t rank = node->rank();
+ shape_data->rank(rank);
+ for (uint32_t index = 0; index < rank; ++index)
+ {
+ if (node->dim(index).known())
+ shape_data->dim(index) = loco::make_dimension(node->dim(index).value());
+ else
+ shape_data->dim(index).unset();
+ }
+ node->annot(std::move(shape_data));
+
+ return true;
}
bool fix_shape(loco::FeatureDecode *node)