bool fix_shape(moco::tf::TFShape *node)
{
- // TODO implement
- throw std::runtime_error("TFShape shape inference NYI");
+ auto shapedata = node->annot<ShapeInferenceData>();
+ if (shapedata != nullptr)
+ {
+ // shape inference is already done for TFShape
+ return false;
+ }
+
+ auto input = node->input();
+ auto input_shape = input->annot<ShapeInferenceData>();
+ if (input_shape == nullptr)
+ {
+ // Input shape is required for TFShape shape inference
+ return false;
+ }
+
+ loco::TensorShape node_shape;
+
+ // Note that input shape becomes node(TFShape)'s value
+ node_shape.rank(1);
+ node_shape.dim(0) = input_shape->rank();
+
+ auto shape_annot = stdex::make_unique<ShapeInferenceData>();
+ shape_annot->tensor_shape(node_shape);
+ node->annot(std::move(shape_annot));
+
+ LOGGER(l);
+ INFO(l) << "Fix TFShape shape = " << node_shape;
+
+ return true;
}
bool fix_shape(moco::tf::TFSqueeze *node)