honor output_dtype in multinominal; cleanup dtype mapping

Author: guschmue

tests/test_backend.py

@@ -175,7 +175,7 @@ def test_trig_ops(self):
     def test_multinomial(self):
         x_val = make_xval([3, 4])
         x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
-        op = tf.multinomial(x, 2)
+        op = tf.multinomial(x, 2, output_dtype=tf.int64)
         output = tf.identity(op, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual, rtol=1e-06)
@@ -726,11 +726,11 @@ def test_cast(self):
 
     @unittest.skip
     def test_onehot(self):
-        # FIXME via onnx-ml ?
+        # no such op in onnx
         x_val = np.array([0, 1, 2], dtype=np.int32)
         depth = 3
         x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
-        x_ = tf.one_hot(x, depth)
+        x_ = tf.one_hot(x, depth, on_value=1, axis=0, off_value=0)
         output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)

tf2onnx/tfonnx.py

@@ -66,20 +66,14 @@ def tensorflow_to_onnx(graph):
         for a in node.node_def.attr:
             attr_cnt[a] += 1
             if a == "dtype":
-                attr[a] = utils.get_tf_dtype(node)
+                attr[a] = utils.map_tf_dtype(node.get_attr("dtype"))
             elif a == "T":
                 dtype = node.get_attr("T")
                 if dtype:
                     if not isinstance(dtype, list):
-                        dtypes[node.name] = utils.TF_TO_ONNX_DTYPE.get(dtype)
-            elif a == "output_type":
-                out_type = node.get_attr("output_type")
-                out_type = utils.TF_TO_ONNX_DTYPE[out_type]
-                attr[a] = out_type
-            elif a == "out_type":
-                out_type = node.get_attr("out_type")
-                out_type = utils.TF_TO_ONNX_DTYPE[out_type]
-                attr[a] = out_type
+                        dtypes[node.name] = utils.map_tf_dtype(dtype)
+            elif a in ["output_type", "output_dtype", "out_type"]:
+                attr[a] = utils.map_tf_dtype(node.get_attr(a))
             elif a == "shape":
                 attr[a] = utils.get_shape(node)
             elif a == "Tperm":
@@ -90,9 +84,7 @@ def tensorflow_to_onnx(graph):
                 onnx_tensor = utils.tf_to_onnx_tensor(node.get_attr(a), name=node.name + ":0")
                 attr[a] = onnx_tensor
             elif a == "DstT":
-                dst = node.get_attr("DstT")
-                dst = tf2onnx.utils.TF_TO_ONNX_DTYPE[dst]
-                attr["to"] = dst
+                attr["to"] =  utils.map_tf_dtype(node.get_attr("DstT"))
             elif a == "SrcT":
                 continue
             elif a in ignored_attr:
@@ -768,12 +760,19 @@ def upsample_op(ctx, node, name, args):
     ctx.remove_input(node, node.input[1])
     return node
 
+
 def multinomial_op(ctx, node, name, args):
     # output_dtype output = Multinomial(T logits, int32 num_samples, @int seed, @int seed2, @type output_dtype)
     sample_size = node.inputs[1].get_tensor_value()
     seed = node.get_attr("seed")
     if seed:
         node.set_attr("seed", float(seed.i))
+    output_dtype = node.get_attr("output_dtype")
+    if output_dtype:
+        output_dtype = output_dtype.i
+    else:
+        output_dtype = onnx_pb.TensorProto.INT32
+    node.set_attr("dtype", output_dtype)
     node.set_attr("sample_size", sample_size[0])
     ctx.remove_input(node, node.input[1])
     return node

tf2onnx/utils.py

@@ -75,7 +75,7 @@
     'min', 'seed', 'ends', 'paddings', 'to', 'gamma', 'width_scale', 'normalize_variance', 'group', 'ratio', 'values',
     'dtype', 'output_shape', 'spatial', 'split', 'input_forget', 'keepdims', 'transA', 'auto_pad', 'border', 'low',
     'linear_before_reset', 'height_scale', 'output_padding', 'shape', 'kernel_shape', 'epsilon', 'size', 'starts',
-    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k'
+    'direction', 'max', 'clip', 'across_channels', 'value', 'strides', 'extra_shape', 'scales', 'k', 'sample_size'
 }
 
 
@@ -148,10 +148,7 @@ def get_shape(node):
         pass
     return dims
 
-
-def get_tf_dtype(node):
-    """Get dtype from tensorflow node."""
-    dtype = node.get_attr("dtype")
+def map_tf_dtype(dtype):
     if dtype:
         dtype = TF_TO_ONNX_DTYPE[dtype]
     return dtype