Support Sub of two variables and Slice in Python ONNX (NeuralNetworkVerification#872)

* expand python onnx parser with slice

* update python onnx parser with slice and sub

* update readme

* Update ONNXParser.py

* add constantOfShape

* change log

Author: wu-haoze

CHANGELOG.md

@@ -6,6 +6,7 @@
   - Fixed bug in the parsing of `transpose` nodes in command line C++ parser.
   - Implemented forward-backward abstract interpretation, symbolic bound tightening, interval arithmetic and simulations for all activation functions.
   - Added the BaBSR heuristic as a new branching strategy for ReLU Splitting
+  - Support Sub of two variables, "Mul" of two constants, Slice, and ConstantOfShape in the python onnx parser
 
 ## Version 2.0.0
 

maraboupy/parsers/ONNXParser.py

@@ -23,7 +23,7 @@
 from onnx import TensorProto
 import itertools
 from copy import copy
-from onnx.reference.ops._op_list import Split_18, Unsqueeze_1
+from onnx.reference.ops._op_list import Split_18, Unsqueeze_1, Slice_10
 
 class ONNXParser:
     """
@@ -167,6 +167,8 @@ def makeMarabouEquations(self, nodeName, makeEquations):
             self.dropout(node)
         elif node.op_type == 'Cast':
             self.cast(node)
+        elif node.op_type == 'ConstantOfShape':
+            self.constantOfShape(node)
         elif node.op_type == 'Reshape':
             self.reshape(node)
         elif node.op_type == 'Flatten':
@@ -177,6 +179,8 @@ def makeMarabouEquations(self, nodeName, makeEquations):
             self.unsqueeze(node)
         elif node.op_type == 'Squeeze':
             self.squeeze(node)
+        elif node.op_type == "Slice":
+            self.slice(node)
         elif node.op_type == "BatchNormalization":
             self.batchNorm(node, makeEquations)
         elif node.op_type == 'Concat':
@@ -282,6 +286,24 @@ def constant(self, node):
                 return
         raise RuntimeError("Could not find value of tensor constant")
 
+    def constantOfShape(self, node):
+        """Function representing a constant tensor of shape
+
+        Args:
+            node (node): ONNX node representing constantOfShape operation
+
+        :meta private:
+        """
+        nodeName = node.output[0]
+        inputName = node.input[0]
+        for attr in node.attribute:
+            if attr.name == "value":
+                value = numpy_helper.to_array(get_attribute_value(attr))
+        assert inputName in self.constantMap
+        shape = self.constantMap[inputName]
+        self.constantMap[nodeName] = np.broadcast_to(value, shape)
+        self.shapeMap[nodeName] = shape
+
     def identity(self, node):
         """Function representing identity
 
@@ -440,6 +462,23 @@ def transpose(self, node):
         elif inputName in self.constantMap:
             self.constantMap[nodeName] = np.transpose(self.constantMap[inputName], perm)
 
+    def slice(self, node):
+        nodeName = node.output[0]
+        inputName = node.input[0]
+        starts = self.constantMap[node.input[1]]
+        ends = self.constantMap[node.input[2]]
+        axes = self.constantMap[node.input[3]]
+        steps = self.constantMap[node.input[4]]
+
+        if inputName in self.varMap:
+            output_data = Slice_10.eval(self.varMap[inputName], starts=starts, ends=ends, axes=axes, steps=steps)
+            self.shapeMap[nodeName] = output_data.shape
+            self.varMap[nodeName] = output_data
+        else:
+            output_data = Slice_10.eval(self.constantMap[inputName], starts=starts, ends=ends, axes=axes, steps=steps)
+            self.shapeMap[nodeName] = output_data.shape
+            self.constantMap[nodeName] = output_data
+
     def unsqueeze(self, node):
         """Function representing unsqueeze
 
@@ -1055,17 +1094,21 @@ def mulEquations(self, node, makeEquations):
             return
 
         multiple = self.constantMap[inputName2]
-        input1 = self.varMap[inputName1]
-        outputVariables = self.makeNewVariables(nodeName)
-        input1 = input1.reshape(-1)
-        outputVariables = outputVariables.reshape(-1)
+        if inputName1 in self.constantMap:
+            input1 = self.constantMap[inputName1]
+            self.constantMap[nodeName] = input1 * multiple
+        else:
+            input1 = self.varMap[inputName1]
+            outputVariables = self.makeNewVariables(nodeName)
+            input1 = input1.reshape(-1)
+            outputVariables = outputVariables.reshape(-1)
 
-        for i in range(len(input1)):
-            e = MarabouUtils.Equation()
-            e.addAddend(multiple, input1[i])
-            e.addAddend(-1, outputVariables[i])
-            e.setScalar(0.0)
-            self.query.addEquation(e)
+            for i in range(len(input1)):
+                e = MarabouUtils.Equation()
+                e.addAddend(multiple, input1[i])
+                e.addAddend(-1, outputVariables[i])
+                e.setScalar(0.0)
+                self.query.addEquation(e)
         return
 
     def addEquations(self, node, makeEquations):
@@ -1245,7 +1288,22 @@ def subEquations(self, node, makeEquations):
         if not makeEquations:
             return
 
-        assert inputName1 in self.varMap and inputName2 in self.constantMap
+        assert inputName1 in self.varMap and (inputName2 in self.constantMap or inputName2 in self.varMap)
+
+        # the difference between the two variables
+        if inputName1 in self.varMap and inputName2 in self.varMap:
+            outputVariables = self.makeNewVariables(nodeName)
+            input1 = self.varMap[inputName1].reshape(-1)
+            input2 = self.varMap[inputName2].reshape(-1)
+            outputVariables = outputVariables.reshape(-1)
+            for i in range(len(input1)):
+                e = MarabouUtils.Equation()
+                e.addAddend(1, input1[i])
+                e.addAddend(-1, input2[i])
+                e.addAddend(-1, outputVariables[i])
+                e.setScalar(0.0)
+                self.query.addEquation(e)
+            return
 
         # Get variables
         inputVars = self.varMap[inputName1].reshape(-1)

maraboupy/test/test_onnx.py

@@ -39,6 +39,18 @@ def test_split_onnx():
     os.remove(presplit_filename)
     os.remove(postsplit_filename)
 
+def test_slice_var():
+    filename =  "slice_test.onnx"
+    evaluateFile(filename)
+
+def test_slice_constant():
+    filename =  "constant_slice_test.onnx"
+    evaluateFile(filename)
+
+def test_sub_var():
+    filename =  "sub_var_test.onnx"
+    evaluateFile(filename)
+
 def test_sub():
     filename =  "test_sub.onnx"
     evaluateFile(filename)

resources/onnx/constant_slice_test.onnx

@@ -0,0 +1,17 @@
+pytorch2.6.0:
+#
+input1
+input2output/Sub"Sub
+main_graphZ
+input1
+

+
+Z
+input2
+

+
+b
+output
+

+
+B