dims can be a list

Author: guschmue

tests/test_backend.py

@@ -184,15 +184,20 @@ def test_expand_dims_known_rank(self):
     def test_expand_dims_one_unknown_rank(self):
         x_val = make_xval([3, 4])
         def func(x):
-            # FIXME: this was tf_placeholder(tf.float32, shape=[None, 4], name=_TFINPUT)
             op = tf.expand_dims(x, 0)
             return tf.identity(op, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    def test_expand_dims_with_list(self):
+        x_val = make_xval([3, 4])
+        def func(x):
+            op = tf.expand_dims(x, [0])
+            return tf.identity(op, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     def _test_expand_dims_more_unknown_rank(self, idx):
         x_val = make_xval([3, 4])
         def func(x):
-            # FIXME: this was tf_placeholder(tf.float32, shape=[None, None], name=_TFINPUT)
             op = tf.expand_dims(x, idx)
             return tf.identity(op, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})

tf2onnx/onnx_opset/tensor.py

@@ -614,7 +614,6 @@ def version_7(cls, ctx, node, **kwargs):
         if dim_node.is_const():
             node.type = "Unsqueeze"
             dim = dim_node.get_tensor_value()
-            # TODO: isn't this always a list ?
             if isinstance(dim, list):
                 dim = dim[0]
             if dim < 0:
@@ -631,6 +630,9 @@ def version_11(cls, ctx, node, **kwargs):
         if dim_node.is_const():
             node.type = "Unsqueeze"
             dim = dim_node.get_tensor_value()
+            if isinstance(dim, list):
+                # tf.expanddims() wants a scalar per doc but quietly accepts a list too.
+                dim = dim[0]
             node.set_attr("axes", [dim])
             ctx.remove_input(node, node.input[1])
             return