Add support for one hot encoding of 0-rank tensor

Signed-off-by: Dagnas <[email protected]>

Author: LoicDagnasSinequa

tests/test_backend.py

@@ -2421,6 +2421,17 @@ def func(x):
                 graph = self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
                 self.assertTrue(len(group_nodes_by_type(graph)["OneHot"]) == 1, "onnx onehot should be used")
 
+    @check_opset_min_version(9, "onehot")
+    def test_onehot_rank0(self):
+        depth = 5
+        for np_dtype in [np.int32, np.int64]:
+            x_val = np.array(3, dtype=np_dtype)
+            for axis in [-1, 0]:
+                def func(x):
+                    x_ = tf.one_hot(x, depth, axis=axis)
+                    return tf.identity(x_, name=_TFOUTPUT)
+                self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     @skip_caffe2_backend("issue undefined dim 1")
     @check_tf_max_version("1.15", "not supported in tf-2.0")
     def test_flatten0(self):

tf2onnx/onnx_opset/tensor.py

@@ -1344,6 +1344,19 @@ def any_version_after9(cls, opset, ctx, node, **kwargs):
         off_on_value = ctx.make_node("Concat", [off_value, on_value], attr={"axis": 0}).output[0]
 
         indices = node.input[0]
+        indices_rank = ctx.get_rank(indices)
+
+        # Add a special support for 0-rank indices, to do so we have to expand the dimension to 1
+        # before the one hot encoding and remove it after.
+        if indices_rank == 0:
+            dims = ctx.make_const(name=utils.make_name('dims'), np_val=np.array([1], dtype=np.int64))
+            indices = ctx.make_node("Expand", [indices, dims.name]).output[0]
+
+            # Axis 0 is supported by TensorFlow for the one-hot encoding of a 0-rank tensor. It should behave
+            # as if axis has been set to -1 so we artificially set it as is here.
+            if node.get_attr('axis').i == 0:
+                node.set_attr('axis', -1)
+
         if ctx.is_target(constants.TARGET_RS6) \
                 and ctx.get_dtype(indices) != onnx_pb.TensorProto.INT64:
             indices = ctx.make_node("Cast", [indices], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
@@ -1367,6 +1380,26 @@ def any_version_after9(cls, opset, ctx, node, **kwargs):
             ctx.set_dtype(new_node.output[0], output_dtype)
             ctx.set_shape(new_node.output[0], ctx.get_shape(node.output[0]))
 
+        # Remove the dimension artificially added in order to support 0-rank indices
+        if indices_rank == 0:
+            nodes = [node]
+            name = utils.make_name(node.name)
+            shape = ctx.get_shape(node.output[0])
+            dtype = ctx.get_dtype(node.output[0])
+            squeeze_node = GraphBuilder(ctx).make_squeeze(
+                {
+                    "axes": [0],
+                    'data': node.output[0]
+                },
+                name=name,
+                dtypes=[dtype],
+                shapes=[shape],
+                return_node=True)
+            ctx.insert_node_on_output(squeeze_node)
+
+            nodes.append(squeeze_node)
+            ctx.update_node_shape_dtype(node, override=True)
+
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
         cls.any_version_after9(9, ctx, node, **kwargs)