Check unsupported dtype for OneHot

Author: winnietsang

onnx_tf/handlers/backend/onehot.py

@@ -1,6 +1,7 @@
 import copy
 import tensorflow as tf
 
+from onnx_tf.common import exception
 from onnx_tf.handlers.backend_handler import BackendHandler
 from onnx_tf.handlers.handler import onnx_op
 from onnx_tf.handlers.handler import tf_func
@@ -10,19 +11,34 @@
 @tf_func(tf.one_hot)
 class OneHot(BackendHandler):
 
+  @classmethod
+  def args_check(cls, node, **kwargs):
+    tensor_dict = kwargs["tensor_dict"]
+    indices = tensor_dict[node.inputs[0]]
+    depth = tensor_dict[node.inputs[1]]
+    if indices.dtype not in [tf.uint8, tf.int32, tf.int64]:
+      exception.OP_UNSUPPORTED_EXCEPT(
+          "OneHot indices must be in uint8 or int32 or int64 " +
+          "but it is currently in " + str(indices.dtype) + " which",
+          "Tensorflow")
+    if depth.dtype not in [tf.int32]:
+      exception.OP_UNSUPPORTED_EXCEPT(
+          "OneHot depth must be in int32 but it is currently in " + str(
+              depth.dtype) + " which", "Tensorflow")
+
   @classmethod
   def version_9(cls, node, **kwargs):
     attrs = copy.deepcopy(node.attrs)
     tensor_dict = kwargs["tensor_dict"]
     indices = tensor_dict[node.inputs[0]]
-    depth = tensor_dict[node.inputs[1]][0]
+    depth = tensor_dict[node.inputs[1]]
     off_value = tensor_dict[node.inputs[2]][0]
     on_value = tensor_dict[node.inputs[2]][1]
     attrs["dtype"] = on_value.dtype
     return [
         cls.make_tensor_from_onnx_node(
             node,
-            inputs=[indices, depth, on_value, off_value],
+            inputs=[indices, depth[0], on_value, off_value],
             attrs=attrs,
             **kwargs)
     ]

onnx_tf/handlers/frontend/onehot.py

@@ -27,7 +27,6 @@ def version_9(cls, node, **kwargs):
     depth = node.inputs[1]
     axis = node.attr.get('axis', -1)
 
-    import pdb; pdb.set_trace()
     on_value = kwargs['consts'][node.inputs[2]].item(0)
     off_value = kwargs['consts'][node.inputs[3]].item(0)
     values = np.array([off_value, on_value])

test/backend/test_onnx_backend.py

@@ -37,6 +37,10 @@
 backend_test.exclude(r'test_mod_[a-z,_]*uint[0-9]+')
 backend_test.exclude(r'test_mod_[a-z,_]*int(8|(16))+')
 
+# TF only support uint8, int32, int64 for indices and int32 for depth in
+# tf.one_hot
+backend_test.exclude(r'test_onehot_[a-z,_]*')
+
 if legacy_opset_pre_ver(7):
   backend_test.exclude(r'[a-z,_]*Upsample[a-z,_]*')