Support for inf or -inf

PINTO0309 · Jun 10, 2022 · 26afae5 · 26afae5
1 parent 3e15658
commit 26afae5
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/sog4onnx/__init__.py b/sog4onnx/__init__.py
@@ -1,3 +1,3 @@
 from sog4onnx.onnx_operation_generator import generate, main
 
-__version__ = '1.0.12'
+__version__ = '1.0.13'
diff --git a/sog4onnx/onnx_operation_generator.py b/sog4onnx/onnx_operation_generator.py
@@ -176,7 +176,7 @@ def generate(
             value_info = onnx.helper.make_tensor_value_info(
                 constant_name,
                 dtype,
-                attr_values.shape
+                attr_values.shape,
             )
             node = onnx.helper.make_node(
                 op_type,
@@ -362,10 +362,39 @@ def main():
             attr_type = attribute[1]
             if attr_type == 'string':
                 attr_type = 'str'
-            if attr_type != 'str':
-                attr_value = ast.literal_eval(attribute[2])
-            else:
+            if attr_type == 'str':
                 attr_value = attribute[2]
+            else:
+                if ('-inf' in attribute[2].lower()) or ('-infinity' in attribute[2].lower()):
+                    lower_attr = attribute[2].lower()
+                    inf_count = lower_attr.count('-inf')
+                    infinity_count = lower_attr.count('-infinity')
+                    if (inf_count + infinity_count) > 1:
+                        print(
+                            f'{Color.RED}ERROR:{Color.RESET} '+
+                            f'Values containing "inf" or "-inf" can only be 1D tensors. \n'+
+                            f'e.g. [inf] or [-inf]'
+                        )
+                        sys.exit(1)
+                    else:
+                        attr_value = [-np.inf]
+
+                elif ('inf' in attribute[2].lower()) or ('infinity' in attribute[2].lower()):
+                    lower_attr = attribute[2].lower()
+                    inf_count = lower_attr.count('inf')
+                    infinity_count = lower_attr.count('infinity')
+                    if (inf_count + infinity_count) > 1:
+                        print(
+                            f'{Color.RED}ERROR:{Color.RESET} '+
+                            f'Values containing "inf" or "-inf" can only be 1D tensors. \n'+
+                            f'e.g. [inf] or [-inf]'
+                        )
+                        sys.exit(1)
+                    else:
+                        attr_value = [np.inf]
+
+                else:
+                    attr_value = ast.literal_eval(attribute[2])
 
             # dtype check
             if attr_type not in AVAILABLE_DTYPES: