support for onehost, support for reshape on int tensors

Author: guschmue

tests/test_backend.py

@@ -805,7 +805,7 @@ def test_onehot0(self):
         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, on_value=5.0, axis=1, off_value=1.0, dtype=tf.float32)
+        x_ = tf.one_hot(x, depth, on_value=5.0, axis=0, off_value=1.0, dtype=tf.float32)
         output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
@@ -821,6 +821,16 @@ def test_onehot1(self):
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
+    def test_onehot2(self):
+        # no such op in onnx
+        x_val = np.array([0, 1, 2, 1, 2, 0, 1, 2, 1, 2], dtype=np.int32)
+        depth = 20
+        x = tf.placeholder(tf.int32, x_val.shape, name=_TFINPUT)
+        x_ = tf.one_hot(x, depth, on_value=5.0, axis=-1, off_value=1.0, dtype=tf.float32)
+        output = tf.identity(x_, name=_TFOUTPUT)
+        actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
+        self.assertAllClose(expected, actual)
+
     @unittest.skipIf(BACKEND in ["caffe2"], "issue undefined dim 1")
     def test_flatten0(self):
         x_val = np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], dtype=np.float32)

tf2onnx/tfonnx.py

@@ -892,12 +892,14 @@ def onehot_op(ctx, node, name, args):
     # until there is no onehot op in onnx, a workaround using gather from eye
     data = node.input[0]
     shape = ctx.get_shape(data)
+    shapeo = ctx.get_shape(node.output[0])
     if len(shape) != 1:
         # TODO: this works for rank=1 but tensorflow supports more than this.
         # Same principle should work but we need to implemtn our own eye.
         raise ValueError("onehot op: only rank1 is supported")
     axis = node.get_attr("axis")
     node.set_attr("axis", axis.i)
+    node.set_attr("axis", 0)
     depth = node.inputs[1].get_tensor_value()[0]
     on = node.inputs[2].get_tensor_value()[0]
     off = node.inputs[3].get_tensor_value()[0]