From 11a0744cd33e217e7c307da50490d3dec43ace7e Mon Sep 17 00:00:00 2001 From: Johan Gunnarsson Date: Tue, 4 Feb 2025 10:02:30 +0100 Subject: [PATCH] MLBEDSW-10181: Infer OFM shape if it's missing There are cases where, for example, a QUANTIZE op has a constant tensor input, but shapless output. With this patch, OFM will inherit IFM's shape for certain ops. Signed-off-by: Johan Gunnarsson Change-Id: Ibcdc503ed56cbe0822e3d5d33dfde76f5620cd07 --- ethosu/regor/tflite/tflite_reader.cpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/ethosu/regor/tflite/tflite_reader.cpp b/ethosu/regor/tflite/tflite_reader.cpp index 50d96324..e5ad5f34 100644 --- a/ethosu/regor/tflite/tflite_reader.cpp +++ b/ethosu/regor/tflite/tflite_reader.cpp @@ -275,6 +275,23 @@ void TfLiteReader::LoadGraphs(const uint8_t *input, const tflite::Model *model, for ( const int tensor_index : *tflite_outputs ) { const auto &ofm = tensors.at(tensor_index); + if ( !ofm->StorageShape() ) + { + // Try to figure out the OFM shape if the OFM shape is unknown + if ( IsUnaryElementwise(op_type) || op_type == OpType::Quantize ) + { + auto ifm = operation->IFM(0); + assert(ifm); + ofm->SetStorageShape(ifm->StorageShape()); + } + else if ( IsBinaryElementwise(op_type) ) + { + auto ifm0 = operation->IFM(0); + auto ifm1 = operation->IFM(1); + assert(ifm0 && ifm1); + ofm->SetStorageShape(Shape::Max(ifm0->StorageShape(), ifm1->StorageShape())); + } + } shapelessTensors = shapelessTensors || !ofm->StorageShape(); assert(tensorQuantization.count(ofm->Uid()) > 0); operation->ConnectOutput(MakeTensorUsage(TensorUsage::OFM, ofm_count++), ofm).Set(tensorQuantization[ofm->Uid()]); -- GitLab