From e4d91dd6184efb1a91cb850a973518512233992d Mon Sep 17 00:00:00 2001
From: Ruediger Ehlers <mail2022@ruediger-ehlers.de>
Date: Tue, 13 Feb 2024 18:17:10 +0100
Subject: [PATCH] Two fixes (but not following the local formatting standards)

- Using a non-two-D Convolution Layer now gives an error message
- If the layer sizes are stored in INT64 instead of INT32, this now works
---
 src/input_parsers/OnnxParser.cpp | 33 ++++++++++++++++++++++++++------
 1 file changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/input_parsers/OnnxParser.cpp b/src/input_parsers/OnnxParser.cpp
index fef37a027c..5026ad2677 100644
--- a/src/input_parsers/OnnxParser.cpp
+++ b/src/input_parsers/OnnxParser.cpp
@@ -411,12 +411,26 @@ Vector<int> getTensorIntValues( const onnx::TensorProto& tensor, const TensorSha
     {
         checkEndianness();
         const char* bytes = raw_data.c_str();
-        const int* ints = reinterpret_cast<const int*>( bytes );
-        for ( int i = 0; i < size; i++ )
-        {
-            int value = *(ints + i);
-            result.append( value );
-        }
+        onnx::TensorProto_DataType dataType = static_cast<onnx::TensorProto_DataType>( tensor.data_type() );
+        
+        if (dataType == onnx::TensorProto_DataType_INT32) {
+    		const int32_t* ints = reinterpret_cast<const int32_t*>( bytes );
+		for ( int i = 0; i < size; i++ )
+		{
+		    int value = *(ints + i);
+		    result.append( value );
+		}
+	} else if (dataType == onnx::TensorProto_DataType_INT64) {
+		const int64_t* ints = reinterpret_cast<const int64_t*>( bytes );
+		for ( int i = 0; i < size; i++ )
+		{
+		    int value = *(ints + i);
+		    result.append( value );
+		}
+	} else {
+		String errorMessage = Stringf( "Illegal data type for integer tensors used in the model. Only INT32 and INT64 supported." );
+            throw MarabouError( MarabouError::ONNX_PARSER_ERROR, errorMessage.ascii() );
+	}
     }
     else
     {
@@ -1239,6 +1253,13 @@ void OnnxParser::convEquations( onnx::NodeProto& node, [[maybe_unused]] bool mak
     // First input should be variable tensor
     String inputNodeName = node.input()[0];
     TensorShape inputShape = _shapeMap[inputNodeName];
+
+    // Added: Check if convolutional layers are only 2D
+    if (inputShape.size()!=4) {
+    	String errorMessage = Stringf( "Onnx '%s' operation has an unsupported number of dimensions -- only 2D convolutions are supported.", node.op_type().c_str() ) ;
+        throw MarabouError( MarabouError::ONNX_PARSER_ERROR, errorMessage.ascii() ) ;
+    }
+
     [[maybe_unused]] unsigned int inputChannels = inputShape[1];
     unsigned int inputWidth = inputShape[2];
     unsigned int inputHeight = inputShape[3];