Fix reduction ERROR_IF cases
Signed-off-by: Kevin Cheng <kevin.cheng@arm.com>
Change-Id: Id0e4ec849a9cf94c9fb04ca999738cc164dbb669
diff --git a/reference_model/src/ops/reduction.cc b/reference_model/src/ops/reduction.cc
index 107c7a8..8c1c4d0 100644
--- a/reference_model/src/ops/reduction.cc
+++ b/reference_model/src/ops/reduction.cc
@@ -50,20 +50,30 @@
if (attribute->axis() < 0 || attribute->axis() >= inputs[0]->getRank())
{
- printNodeValidationError("Reduce axis must between [0, input_rank - 1]");
+ printNodeValidationError("ReduceOp: axis must between [0, input_rank - 1]");
return 1;
}
- if (inputs[0]->matchRank(*outputs[0]))
+ if (inputs[0]->matchRankType(*outputs[0]))
{
- printNodeValidationError("Input and output tensor ranks must match");
+ printNodeValidationError("ReduceOp: Input and output tensor ranks must match");
+ return 1;
+ }
+
+ if (outputs[0]->getShape()[attribute->axis()] != 1)
+ {
+ printNodeValidationError("ReduceOp: Output tensor shape[axis] needs to be 1.");
return 1;
}
in = dynamic_cast<TosaReference::TensorTemplate<TIn>*>(inputs[0]);
out = dynamic_cast<TosaReference::TensorTemplate<TOut>*>(outputs[0]);
- ASSERT_MEM(in && out);
+ if ((!in) || (!out))
+ {
+ printNodeValidationError("ReduceOp: Input or output fail to cast to Eigen tensor since rank/type not expected");
+ return 1;
+ }
dims[0] = this->attribute->axis();