Reactivate constant folding (for ConcatV2)

Allow output shape of (?,)
Fix Squeeze op when tensor shape is (?,)

Author: jignparm

tf2onnx/graph.py

@@ -1067,7 +1067,9 @@ def make_onnx_graph_io(self, ids):
             shape = self.get_shape(name)
 
             utils.make_sure(dtype is not None, "missing output dtype for " + name)
-            utils.make_sure(shape is not None, "missing output shape for " + name)
+            # TODO: allow None output shape or not? e.g. shape=(?,)
+            #utils.make_sure(shape is not None, "missing output shape for " + name)
+            if shape is None: logger.warning(f"missing output shape for " + name)
 
             v = utils.make_onnx_inputs_outputs(name, dtype, shape)
             tensor_value_infos.append(v)

tf2onnx/onnx_opset/tensor.py

@@ -196,6 +196,7 @@ def version_1(cls, ctx, node, **kwargs):
             shape = ctx.get_shape(node.input[0])
             utils.make_sure(shape is not None, "squeeze input shape cannot be None")
             axis = [i for i, j in enumerate(shape) if j == 1]
+            if not axis: axis = [0]
         node.set_attr("axes", axis)
 
     @classmethod

tf2onnx/tfonnx.py

@@ -42,17 +42,17 @@ def rewrite_constant_fold(g, ops):
     tensorflow missed something, make another pass over the graph and fix want we care about.
     """
     func_map = {
-        "Add": np.add,
-        "GreaterEqual": np.greater_equal,
-        "Cast": np.cast,
+        # "Add": np.add,
+        # "GreaterEqual": np.greater_equal,
+        # "Cast": np.cast,
         "ConcatV2": np.concatenate,
-        "Less": np.less,
-        "ListDiff": np.setdiff1d,
-        "Mul": np.multiply,
-        "Pack": np.stack,
-        "Range": np.arange,
-        "Sqrt": np.sqrt,
-        "Sub": np.subtract,
+        # "Less": np.less,
+        # "ListDiff": np.setdiff1d,
+        # "Mul": np.multiply,
+        # "Pack": np.stack,
+        # "Range": np.arange,
+        # "Sqrt": np.sqrt,
+        # "Sub": np.subtract,
     }
     ref_cnt_per_node = {}
     for idx, op in enumerate(ops):
@@ -466,7 +466,7 @@ def compat_handler(ctx, node, **kwargs):
                  rewrite_single_direction_lstm, rewrite_bi_direction_lstm,
                  rewrite_single_direction_gru, rewrite_bi_direction_gru,
                  rewrite_custom_rnn_cell, rewrite_generic_loop, rewrite_cond,
-                 rewrite_biasadd_with_conv2d,
+                 rewrite_biasadd_with_conv2d, rewrite_constant_fold
                  ]
 
     if custom_rewriter is not None: