From ec5586c198e81fba43f598af9ecd7a54cf460ea3 Mon Sep 17 00:00:00 2001
From: Kevin Cheng <kevin.cheng@arm.com>
Date: Wed, 6 Oct 2021 14:37:37 -0700
Subject: Fix reduction ERROR_IF cases

Signed-off-by: Kevin Cheng <kevin.cheng@arm.com>
Change-Id: Id0e4ec849a9cf94c9fb04ca999738cc164dbb669
---
 reference_model/src/ops/reduction.cc | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

(limited to 'reference_model/src/ops')

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 @@ int ReduceNode<Rank, Dtype>::checkTensorAttributes()
 
     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();
 
-- 
cgit v1.2.1